diff --git a/CLAUDE.md b/CLAUDE.md index 019a0961763..3f0e86bc50f 100644 --- a/CLAUDE.md +++ b/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`. | diff --git a/benchmarks/comm/bench_pcie_ipc_all_reduce.py b/benchmarks/comm/bench_pcie_ipc_all_reduce.py new file mode 100644 index 00000000000..7a5296fed60 --- /dev/null +++ b/benchmarks/comm/bench_pcie_ipc_all_reduce.py @@ -0,0 +1,686 @@ +# 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. + +"""PCIe IPC all-reduce latency, against NCCL. + +Intended for intra-node PCIe machines without NVLink, which is what the kernel +is tuned for. NCCL is the baseline because it is what the caller falls back to +on a shape the kernels cannot run. + +One run covers ONE world size. The world size comes from the launcher and the +collective uses ``dist.group.WORLD`` throughout -- there is no subgroup logic -- +so a 2/4/8-rank comparison needs three launches: + + sudo nvidia-smi -lgc 2520,2520 # nothing below pins clocks; boost drift + # swamps the 2-3% differences being measured + + for n in 2 4 8; do + torchrun --standalone --nproc_per_node=$n \\ + benchmarks/comm/bench_pcie_ipc_all_reduce.py --json bench_tp$n.json + done + +Hidden size follows the world size unless overridden, so no other flag is +needed. The defaults are a sweep choice, not a restriction -- --hidden takes +any size the kernels admit. + +To measure what a protocol fix costs, rather than how it compares to NCCL: + + for n in 2 4 8; do + timeout 3000 torchrun --standalone --nproc_per_node=$n \\ + benchmarks/comm/bench_pcie_ipc_all_reduce.py \\ + --protocol-ab auto --json ab_tp$n.json + done + +The external timeout is not boilerplate. At 8 ranks the historically faithful +baseline for the staged kernels IS the cross-island protocol that was removed +for being broken, so rank skew can leave it spinning with nothing to time it +out from inside. + +'auto' compares each shape against the protocol it actually used to run, which +at 8 ranks differs between the pack and staged kernels. Naming a switch +explicitly is still allowed, and rows it does not historically fit are marked +SYNTHETIC. + +To measure the launch configuration rather than take the policy's seed default: + + torchrun --standalone --nproc_per_node=8 \\ + benchmarks/comm/bench_pcie_ipc_all_reduce.py --tune --json bench_tuned.json + +Rows whose configuration tuning changed are annotated with what the seed would +have chosen. The result is persisted, so a later run without --tune reuses it. +Without a persisted result the library warns once that it is running seed +configurations, which are workable rather than fast. + +Options: + --hidden N Hidden size (default: 6144 at 8 ranks, 4096 at 4, 2048 at 2) + --batches a,b,c Batch sizes to sweep + --dtype bfloat16 (default) or float16 + --json FILE Write results to JSON + --tune Measure the launch configuration instead of taking the seed + --tune-cache FILE Where tuned configurations are read from and written to +""" + +import argparse +import json +import os +from typing import Dict, List + +import torch +import torch.distributed as dist + +import flashinfer.comm as comm +from flashinfer.comm import pcie_ipc_ar +from flashinfer.comm.pcie_ipc_policy import IpcVariant +from flashinfer.jit.comm import gen_pcie_ipc_comm_debug_module +from flashinfer.testing.utils import bench_gpu_time + +_DEFAULT_HIDDEN = {2: 2048, 4: 4096, 8: 6144} +_DEFAULT_BATCHES = [1, 2, 4, 8, 16, 32, 64, 128] + +# Pin the iteration counts. Left to auto-tune, each rank derives its own count +# from its own timings, and a collective where the ranks disagree on how many +# times to call it deadlocks. +_BENCH_KWARGS = dict( + use_cuda_graph=True, + # Each replay carries a fixed cost -- the launch plus the event pair timing + # it -- that the reported number amortises over the iterations in the graph, + # so anything compared against these numbers must use the same count. + num_iters_within_graph=20, + dry_run_iters=5, + repeat_iters=20, + # The timed callables close over their tensors rather than taking them as + # arguments, which is all the cold-L2 helper inspects. + cold_l2_cache=False, +) + + +def _median_us(samples) -> float: + # bench_gpu_time reports milliseconds. + return float(torch.tensor(samples).median()) * 1e3 + + +def _group_median_us(samples, device, group) -> float: + """Median over iterations of the group maximum at each iteration. + + Order matters here. Taking each rank's median and then the max across ranks + computes ``max_rank median_iter``, which is not what a collective costs: if + rank 0 is slow on odd iterations and rank 4 on even ones, every rank's own + median is low while every actual iteration had a straggler. The tail is + ``median_iter max_rank``, so the max has to be taken per sample, before the + median. + + Requires every rank to have the same samples in the same order, which is why + the iteration counts are pinned rather than auto-tuned; asserted below + because silently reducing across mismatched vectors would produce a number + that looks fine. + """ + local = list(samples) + # One collective for both bounds: max of n, and max of -n (i.e. -min). + bounds = torch.tensor([len(local), -len(local)], dtype=torch.int64, device=device) + dist.all_reduce(bounds, op=dist.ReduceOp.MAX, group=group) + if int(bounds[0].item()) != -int(bounds[1].item()): + raise RuntimeError( + "ranks produced different sample counts " + f"({-int(bounds[1].item())}..{int(bounds[0].item())}); " + "per-iteration aggregation needs them aligned" + ) + t = torch.as_tensor(local, dtype=torch.float64, device=device) + dist.all_reduce(t, op=dist.ReduceOp.MAX, group=group) + # bench_gpu_time reports milliseconds. + return float(t.median()) * 1e3 + + +def _group_all(flag: bool, device, group) -> bool: + """True only when every rank says True. + + Correctness cannot be judged locally here. The cross-island race this + harness can rebuild produced errors confined to *one island*: rank 0 came + out clean while ranks 4-7 had millions of wrong elements. A rank-local check + would have published a cost for a baseline that was wrong. + """ + t = torch.tensor([1 if flag else 0], dtype=torch.int32, device=device) + dist.all_reduce(t, op=dist.ReduceOp.MIN, group=group) + return bool(t.item()) + + +def _provenance(config, seed, tuning: bool) -> str: + """How this row's configuration was chosen.""" + if seed != config: + return f" (seed: {seed})" if tuning else "" + return " (seed -- no tuned entry)" if not tuning else "" + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--hidden", type=int, default=None) + p.add_argument("--batches", type=str, default=None) + p.add_argument("--dtype", choices=["bfloat16", "float16"], default="bfloat16") + p.add_argument("--json", type=str, default=None) + p.add_argument( + "--tune", + action="store_true", + help=( + "Measure the launch configuration instead of taking the policy " + "seed, then report against NCCL. Collective and slow (a few " + "seconds per batch bucket); pin clocks first." + ), + ) + p.add_argument( + "--tune-cache", + type=str, + default=None, + help=( + "Where tuned configurations are read from and, under --tune, " + "written to. Defaults to the workspace dir. Reading happens " + "with or without --tune, so a later run reuses the result." + ), + ) + p.add_argument( + "--protocol-ab", + choices=[ + "auto", + "per-block-epoch", + "no-block-epoch", + "no-barrier-entry-sync", + ], + default=None, + help=( + "Instead of comparing to NCCL, measure what one protocol mechanism " + "costs: rebuild the kernels with it disabled and time both in the " + "same process, A-B-A per shape, reporting the A/A spread as a noise " + "floor. 'auto' compares each shape against the mechanism it " + "actually used to run, which at 8 ranks differs between the pack " + "and staged kernels. Every switch produces an incorrect build -- run " + "under an external timeout, since one of them can spin. See the " + "module docstring." + ), + ) + return p.parse_args() + + +def _historical_switch(world_size, config): + """Which A/B switch actually rebuilds what this shape used to run. + + Keys on the kernel ``config`` names -- ``(world_size, config.variant)`` + alone -- not on the shape that reached it. The kernels do not share one + history: TP2 and TP4 always double-buffered by per-block parity, so + ``per-block-epoch`` restores their past. The two topology-staged TP8 kernels + never had an epoch at all -- their scratch was pinned to half 0 and reused + across calls, which is what ``no-block-epoch`` rebuilds. The TP8 pack kernel + did have per-block parity, like TP2/TP4, and so does ``FLAT_STAGED``: it is + the same generic template TP4 runs. + + Running the other switch is not wrong, it just measures a protocol that + never shipped, so rows are labelled rather than refused. + """ + tp8_topo_staged = world_size == 8 and config.variant in ( + IpcVariant.STAGED, + IpcVariant.STAGED_RING, + ) + return "no-block-epoch" if tp8_topo_staged else "per-block-epoch" + + +# The two switches that rebuild an *epoch* protocol. Which of them is faithful +# depends on the kernel; every other switch rebuilds a state that no kernel ever +# had a different version of, so it is historical everywhere. +_EPOCH_SWITCHES = ("per-block-epoch", "no-block-epoch") + + +def _is_historical(world_size, config, switch): + if switch not in _EPOCH_SWITCHES: + return True + return switch == _historical_switch(world_size, config) + + +def _protocol_ab_plan(world_size, shape_config, mode): + """Which batches each A/B leg should actually run. + + ``shape_config`` is ``{batch: IpcLaunchConfig}`` for the configurations this + run will really launch -- read off the workspace, not looked up anywhere -- + so the split follows the kernels being benchmarked. + + Returns ``{switch: [batch, ...]}``. In ``auto`` this partitions the shapes by + the protocol they used to run, and that partition is the executable plan, not + just a labelling rule: a leg that runs a shape it does not own executes a + protocol that shape never had, and one of those -- the fixed-half build on a + kernel that always double-buffered -- wedges its sentinel loop outright. A + row discarded at reporting time has already run by then. + + How many legs come back depends on the sweep. One leg is the normal result, + not a degenerate one: world sizes 2 and 4 have a single history, so only a + world-size-8 sweep whose shapes reach both the pack and the topology-staged + kernels partitions into two. + + With an explicit switch the caller has asked for exactly that comparison, so + every admitted shape runs and the mismatched ones are reported as SYNTHETIC. + """ + if mode != "auto": + return {mode: sorted(shape_config)} if shape_config else {} + plan = {} + for batch, config in shape_config.items(): + plan.setdefault(_historical_switch(world_size, config), []).append(batch) + # Unsafe legs last. 'no-block-epoch' rebuilds a protocol that can spin + # forever, and a hang there would otherwise take down the safe leg's results + # with it -- results that were already complete and cost nothing to keep. + order = sorted(plan, key=lambda sw: (sw == "no-block-epoch", sw)) + return {switch: sorted(plan[switch]) for switch in order} + + +def _run_ab_legs(plan, run_sweep, on_switch_done=None): + """Drive A-B-A for every switch in `plan`, over that switch's shapes only. + + Split out from the timing so it can be tested without a GPU. The property + that matters is not which switch a shape is *labelled* with but which leg + actually executes it: a leg running a shape it does not own runs a protocol + that shape never had, and one of those wedges outright. That mistake was + made once already, with a correct labelling function sitting right next to + it, so the wiring is now a seam a test can hold. + + ``run_sweep(broken_switch, batches) -> {batch: result}``; ``broken_switch`` + is None for the shipping build. + """ + legs = {} + for switch, leg_batches in plan.items(): + legs[switch] = ( + run_sweep(switch, leg_batches), + run_sweep(None, leg_batches), + run_sweep(switch, leg_batches), + ) + if on_switch_done is not None: + on_switch_done(switch, legs[switch]) + return legs + + +def _time_one_batch(workspace, batch, hidden, dtype, device, group): + """Time one shape, after proving this build still computes it correctly. + + The correctness check is not decoration. Two of the three A/B baselines are + protocols known to be broken, and a broken protocol that has landed on the + wrong half but not yet wedged its spin loop still produces a perfectly + timeable kernel. Without this, such a run reports a number that means + nothing. + """ + ref_src = torch.randint(1, 16, (batch, hidden), dtype=torch.int32, device=device) + inp = ref_src.to(dtype) + ref = inp.clone() + dist.all_reduce(ref, group=group) + + config = workspace.launch_config(inp) + if config is None: + return None + dst = torch.empty_like(inp) + torch.cuda.synchronize() + workspace.rebind_stream() + + workspace.all_reduce(inp, out=dst, config=config) + torch.cuda.synchronize() + correct = _group_all(bool(torch.equal(dst, ref)), device, group) + if not correct: + # Nothing to learn from timing a build that is already wrong, and this + # is exactly where the unsafe baselines fail -- running the graph + # hundreds more times mostly buys a chance to spin forever instead of + # reporting. Safe to branch: `correct` is group-wide, so every rank + # leaves here together. + return None, None, config, False, False + + # The check above bound the workspace to the current stream, and + # bench_gpu_time warms up on a side stream before capturing. Sequential, not + # concurrent -- the synchronize above is the ordering the escape hatch asks + # for -- so release the binding again. + workspace.rebind_stream() + samples = bench_gpu_time( + lambda: workspace.all_reduce(inp, out=dst, config=config), **_BENCH_KWARGS + ) + us = _group_median_us(samples, device, group) + rank_us = _median_us(samples) + + # And once more afterwards. The pre-check only says this build was right + # once; the timed run then replays the same call hundreds of times inside a + # graph, which is the repetition a broken protocol needs to show itself. It + # still cannot falsify a race -- the payload is constant, so an overwrite may + # store a bit-identical value -- but it catches anything that has actually + # gone wrong by the end, and that costs one comparison. + torch.cuda.synchronize() + workspace.rebind_stream() + workspace.all_reduce(inp, out=dst, config=config) + torch.cuda.synchronize() + post_correct = _group_all(bool(torch.equal(dst, ref)), device, group) + return us, rank_us, config, correct, post_correct + + +def _run_protocol_ab(args, group, world_size, hidden, batches, dtype, device, rank): + """Time the shipping protocol against one with a fix compiled out. + + The workspace resolves its module through a module-global, so swapping that + is enough to put an instrumented build under the same wrapper the shipping + path uses. A-B-A: the two A runs bracket B, and their spread is the noise + floor any claimed delta has to clear. + + **One workspace per batch, per leg.** A shared workspace would let one shape + contaminate the next, and specifically so on the side that matters: the + per-block baseline *is* the protocol whose parity desynchronises for good + once the grid changes, so every batch after the first grid change would be + timed against state left by its predecessors rather than against a clean + single-shape run. It usually does not hang -- the default batch list happens + to avoid the shapes that wedge immediately -- and "did not hang" is not the + same as "measured the thing we named". + + ``auto`` runs one sweep per switch and keeps, for each shape, only the sweep + whose switch matches that shape's actual history. At 8 ranks that is the + only way to get a complete historically faithful set of rows, because the + pack and staged kernels came from different protocols. + """ + # (no_block_epoch, per_block_epoch, no_barrier_entry_sync) + _SWITCH = { + "no-block-epoch": (1, 0, 0), + "per-block-epoch": (0, 1, 0), + "no-barrier-entry-sync": (0, 0, 1), + } + + # Work out which shape maps to which kernel before building anything, so + # auto mode can run only the switches some shape actually needs. Running a + # switch nothing needs is not merely wasted time: 'no-block-epoch' at 2 or 4 + # ranks removes the double buffer from a kernel that has always had one, and + # that wedges the sentinel loop outright. The probe reads the same seed + # configurations the timed legs below will launch, so the plan describes + # what runs. + probe = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * max(batches), dtype=dtype + ) + try: + shape_config = {} + for batch in batches: + cfg = probe.launch_config( + torch.empty(batch, hidden, dtype=dtype, device=device) + ) + if cfg is not None: + shape_config[batch] = cfg + finally: + probe.destroy() + + plan = _protocol_ab_plan(world_size, shape_config, args.protocol_ab) + if not plan: + if rank == 0: + print("no supported shape in this batch list; nothing to compare") + return + + def measure(broken_switch, leg_batches): + """One sweep over `leg_batches`, with `broken_switch` compiled out or not.""" + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + original = pcie_ipc_ar.gen_pcie_ipc_comm_module + flags = _SWITCH[broken_switch] if broken_switch else (0, 0, 0) + + def _gen(): + return gen_pcie_ipc_comm_debug_module(0, 0, *flags) + + pcie_ipc_ar.gen_pcie_ipc_comm_module = _gen + try: + out = {} + for batch in leg_batches: + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * batch, dtype=dtype + ) + try: + got = _time_one_batch(ws, batch, hidden, dtype, device, group) + finally: + ws.destroy() + if got is not None: + out[batch] = got + return out + finally: + pcie_ipc_ar.gen_pcie_ipc_comm_module = original + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + + # A-B-A per switch, over that switch's own shapes only. B runs the same + # subset, so all three legs of a comparison execute an identical workload. + rows: List[Dict[str, object]] = [] + + def _flush(): + """Write what is finished so far, atomically. + + Granularity is one complete A-B-A group, not one leg: without both A + legs there is no A/A spread, so a half-finished switch has nothing worth + keeping. What this does buy is that a switch which spins cannot take an + earlier switch's finished rows down with it -- the unsafe one is ordered + last for the same reason. rename() makes a reader see either the old + file or a complete new one. + """ + if rank != 0 or not args.json: + return + tmp = args.json + ".partial" + with open(tmp, "w") as f: + json.dump(rows, f, indent=2) + os.replace(tmp, args.json) + + def _emit(switch, leg): + """Report one finished leg, then persist.""" + for batch in plan[switch]: + if any(batch not in part for part in leg): + continue + a1, b, a2 = (part[batch] for part in leg) + with_fix, with_fix_rank0, config, fixed_ok, fixed_post = b + baseline_ok = a1[3] and a2[3] and a1[4] and a2[4] + # A leg that failed its pre-check returns no timing at all. + timed = baseline_ok and fixed_ok and fixed_post + without = (a1[0] + a2[0]) / 2 if timed else float("nan") + drift = abs(a1[0] - a2[0]) / without * 100 if timed else float("nan") + historical = _historical_switch(world_size, config) + is_historical = _is_historical(world_size, config, switch) + ok = timed + row = { + "batch": batch, + "hidden": hidden, + "world_size": world_size, + "switch": switch, + "unsafe_baseline": switch + in ("no-block-epoch", "no-barrier-entry-sync"), + "historical_switch": historical + if switch in _EPOCH_SWITCHES + else switch, + "synthetic_baseline": not is_historical, + "blocks": config.blocks, + "threads": config.threads, + "variant": config.variant.name, + "with_fix_us": with_fix, + "with_fix_rank0_us": with_fix_rank0, + "timed": timed, + "latency_agg": "median_of_per_iteration_group_max", + "with_fix_correct": fixed_ok and fixed_post, + "baseline_correct": baseline_ok, + "correctness_is_group_wide": True, + "aa_drift_pct": drift, + } + # A cost against a baseline that computed the wrong answer is not a + # cost, so do not publish one. + if ok: + row["without_fix_us"] = without + row["cost_pct"] = (with_fix - without) / without * 100 + rows.append(row) + if rank == 0: + if not (fixed_ok and fixed_post): + note = "WRONG RESULT" + elif not baseline_ok: + note = "baseline wrong" + else: + note = "historical" if is_historical else "SYNTHETIC" + if switch in ("no-block-epoch", "no-barrier-entry-sync"): + note += " UNSAFE" + if not is_historical: + note += f" (history: {historical})" + if ok: + cost = f"{(with_fix - without) / without * 100:>+7.1f}%" + print( + f"{batch:>7} {config.blocks:>7} {without:>12.2f} " + f"{with_fix:>10.2f} {cost} {drift:>6.1f}% {note}" + ) + else: + print( + f"{batch:>7} {config.blocks:>7} {'--':>12} {'--':>10} " + f"{'--':>8} {'--':>6} {note} (not timed)" + ) + _flush() + + if rank == 0: + label = ( + "auto (per-shape history: " + + ", ".join(f"{sw} x{len(b)}" for sw, b in plan.items()) + + ")" + if args.protocol_ab == "auto" + else args.protocol_ab + ) + print(f"world_size={world_size} hidden={hidden} A/B on {label}") + print( + f"{'batch':>7} {'blocks':>7} {'without(us)':>12} {'with(us)':>10} " + f"{'cost':>8} {'A/A':>7} baseline" + ) + print( + " (latency: median of the per-iteration group max; " + "correctness: group-wide)" + ) + + _run_ab_legs(plan, measure, on_switch_done=_emit) + + if rank == 0: + print( + " historical = the baseline is what this shape actually used to run; " + "SYNTHETIC = a protocol it never ran; UNSAFE = the baseline itself is " + "known-broken, so its timings describe this workload only." + ) + if args.json: + print(f"wrote {args.json} ({len(rows)} rows, flushed after each leg)") + + +def main() -> None: + args = _parse_args() + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group("nccl") + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + + hidden = args.hidden if args.hidden is not None else _DEFAULT_HIDDEN.get(world_size) + if hidden is None: + raise ValueError(f"--hidden is required at {world_size} ranks") + batches = ( + [int(b) for b in args.batches.split(",")] + if args.batches + else list(_DEFAULT_BATCHES) + ) + dtype = getattr(torch, args.dtype) + + if args.protocol_ab is not None: + _run_protocol_ab(args, group, world_size, hidden, batches, dtype, device, rank) + dist.destroy_process_group() + return + + workspace = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=hidden * max(batches), + dtype=dtype, + tune_cache=args.tune_cache, + ) + # Recorded before tuning overwrites it: under --tune this shows what the + # search changed, without it which rows found no cache entry. + seed_configs = { + batch: workspace.launch_config( + torch.empty(batch, hidden, dtype=dtype, device=device) + ) + for batch in batches + } + if args.tune: + if rank == 0: + print(f"tuning {len(batches)} batch buckets at hidden {hidden} ...") + torch.cuda.synchronize() + workspace.rebind_stream() + workspace.tune([hidden], dtype=dtype) + if rank == 0: + print(f"world_size={world_size} hidden={hidden} dtype={args.dtype}") + print(f"profile={workspace.profile} ({workspace.profile_reason})") + print(f"{'batch':>7} {'ours(us)':>10} {'nccl(us)':>10} {'speedup':>9} config") + + rows: List[Dict[str, object]] = [] + for batch in batches: + inp = torch.randn(batch, hidden, dtype=dtype, device=device) + # The tuned answer, so the reported configuration is the one the timed + # call below actually runs. + config = workspace.tuned_launch_config(inp) + # None is a capability answer, never "nobody measured this": an + # unmeasured shape still gets the seed. + if config is None: + if rank == 0: + print( + f"{batch:>7} {'-':>10} {'-':>10} {'-':>9} " + "unsupported, would fall back" + ) + continue + out = torch.empty_like(inp) + + # bench_gpu_time warms up on a side stream before capturing a graph. + # That is sequential, not concurrent, so tell the workspace it may move + # its binding -- after making sure the previous stream is drained. + torch.cuda.synchronize() + workspace.rebind_stream() + # A collective finishes when its slowest rank does, and which rank that + # is can change from iteration to iteration -- so the max is taken per + # sample and the median after, not the other way round. The local + # medians stay in the JSON so an older result set can still be lined up + # against a newer one. + ours_samples = bench_gpu_time( + lambda: workspace.all_reduce(inp, out=out), **_BENCH_KWARGS + ) + nccl_inp = inp.clone() + nccl_samples = bench_gpu_time( + lambda: dist.all_reduce(nccl_inp, group=group), **_BENCH_KWARGS + ) + ours_us = _group_median_us(ours_samples, device, group) + nccl_us = _group_median_us(nccl_samples, device, group) + ours_rank_us = _median_us(ours_samples) + nccl_rank_us = _median_us(nccl_samples) + rows.append( + { + "batch": batch, + "hidden": hidden, + "world_size": world_size, + "dtype": args.dtype, + "ours_us": ours_us, + "nccl_us": nccl_us, + "ours_rank0_us": ours_rank_us, + "nccl_rank0_us": nccl_rank_us, + "latency_agg": "median_of_per_iteration_group_max", + "speedup": nccl_us / ours_us, + "blocks": config.blocks, + "threads": config.threads, + "variant": config.variant.name, + } + ) + if rank == 0: + print( + f"{batch:>7} {ours_us:>10.2f} {nccl_us:>10.2f} " + f"{nccl_us / ours_us:>8.2f}x blocks={config.blocks} " + f"threads={config.threads} variant={config.variant.name}" + + _provenance(config, seed_configs.get(batch), args.tune) + ) + + if rank == 0 and args.json: + with open(args.json, "w") as f: + json.dump(rows, f, indent=2) + print(f"wrote {args.json}") + + workspace.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/csrc/pcie_ipc_all_reduce.cu b/csrc/pcie_ipc_all_reduce.cu new file mode 100644 index 00000000000..bdbf102c88d --- /dev/null +++ b/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 + +#include + +#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(world_size), max_numel, static_cast(elem_size), + static_cast(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 ipc_ptrs, int64_t rank, int64_t max_numel, int64_t elem_size, + int64_t max_blocks) { + const int world_size = static_cast(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(elem_size), + static_cast(max_blocks)); + handle->views = fi::make_peer_views(ptrs, world_size, static_cast(rank), handle->layout); + handle->rank = static_cast(rank); + handle->world_size = world_size; + handle->max_blocks = static_cast(max_blocks); + handle->max_numel = max_numel; + handle->elem_size = static_cast(elem_size); + + cudaError_t err = cudaMemset(reinterpret_cast(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(handle); +} + +void pcie_ipc_dispose(fptr_t handle) { delete reinterpret_cast(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(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(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(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(static_cast(inp.data_ptr()), + static_cast(out.data_ptr()), numel, h->views, + h->rank, h->world_size, h->max_blocks, h->max_numel, + static_cast(blocks), static_cast(threads), algo, + enable_pdl, stream); + break; + case float16_code: + err = fi::all_reduce( + static_cast(inp.data_ptr()), static_cast(out.data_ptr()), numel, + h->views, h->rank, h->world_size, h->max_blocks, h->max_numel, static_cast(blocks), + static_cast(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); diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 103c0e82453..63a83e58bb4 100644 --- a/docs/api/comm.rst +++ b/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 ----------------------------------- diff --git a/flashinfer/aot.py b/flashinfer/aot.py index e49c4567329..292ced9ae54 100644 --- a/flashinfer/aot.py +++ b/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", diff --git a/flashinfer/comm/__init__.py b/flashinfer/comm/__init__.py index b763156bbb1..9b9eaef7b6b 100644 --- a/flashinfer/comm/__init__.py +++ b/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 diff --git a/flashinfer/comm/pcie_ipc_ar.py b/flashinfer/comm/pcie_ipc_ar.py new file mode 100644 index 00000000000..6688e9c5010 --- /dev/null +++ b/flashinfer/comm/pcie_ipc_ar.py @@ -0,0 +1,1002 @@ +""" +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. +""" + +import functools +import hashlib +import os +import warnings +from types import SimpleNamespace +from typing import Dict, List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from ..api_logging import flashinfer_api +from ..trace.templates.comm import pcie_ipc_all_reduce_trace +from ..jit.comm import gen_pcie_ipc_comm_module +from ..utils import register_custom_op +from .cuda_ipc import create_shared_buffer, free_shared_buffer +from .pcie_ipc_policy import IpcLaunchConfig, get_pcie_ipc_launch_config +from .pcie_ipc_topology import resolve_pcie_ipc_profile +from .pcie_ipc_tuning import ( + PCIE_IPC_CUSTOM_OP, + TUNE_BATCHES, + TUNE_REPEAT, + TUNE_WARMUP, + PcieIpcAllReduceRunner, + cache_covers_workspace, + default_cache_path, + pack_config, + pcie_ipc_tuning_config, + resolve_tuned_config, + tuned_batches_for, + warn_no_tune_group, +) + +_SUPPORTED_WORLD_SIZES = (2, 4, 8) +# Mirrors the launcher, which hard-checks a 2-byte element size: the kernels +# address whole 16-byte packs and are instantiated for half and nv_bfloat16 +# only. Rejecting here turns that into an unsupported shape rather than an +# ICHECK partway through a collective. +_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) + + +@functools.cache +def get_pcie_ipc_comm_module(): + module = gen_pcie_ipc_comm_module().build_and_load() + + @register_custom_op("flashinfer::pcie_ipc_workspace_size", mutates_args=[]) + def workspace_size( + world_size: int, max_numel: int, elem_size: int, max_blocks: int + ) -> int: + return module.pcie_ipc_workspace_size( + world_size, max_numel, elem_size, max_blocks + ) + + @register_custom_op("flashinfer::pcie_ipc_init", mutates_args=["ipc_ptrs"]) + def init( + ipc_ptrs: List[int], + rank: int, + max_numel: int, + elem_size: int, + max_blocks: int, + ) -> int: + return module.pcie_ipc_init(ipc_ptrs, rank, max_numel, elem_size, max_blocks) + + @register_custom_op("flashinfer::pcie_ipc_dispose", mutates_args=["handle"]) + def dispose(handle: int) -> None: + module.pcie_ipc_dispose(handle) + + @register_custom_op("flashinfer::pcie_ipc_all_reduce", mutates_args=["out"]) + def all_reduce( + handle: int, + inp: torch.Tensor, + out: torch.Tensor, + blocks: int, + threads: int, + variant: int, + enable_pdl: bool, + ) -> None: + module.pcie_ipc_all_reduce( + handle, inp, out, blocks, threads, variant, enable_pdl + ) + + return SimpleNamespace( + workspace_size=workspace_size, + init=init, + dispose=dispose, + all_reduce=all_reduce, + ) + + +class PcieIpcAllReduceWorkspace: + """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. :meth:`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 + :meth:`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. + + Launch configurations start from a seed default that is workable rather + than fast (see :mod:`~flashinfer.comm.pcie_ipc_policy`). Tune once to + replace it with measurements from this machine; the result is persisted and + later processes pick it up when the workspace is built. Tuning never changes + which shapes are supported, 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 + """ + + def __init__( + self, + group: ProcessGroup, + max_numel: int, + dtype: torch.dtype = torch.bfloat16, + max_blocks: int = 128, + profile: Optional[str] = None, + tune_batches: Sequence[int] = TUNE_BATCHES, + tune_cache: Optional[str] = None, + ) -> None: + # Construction is a staged transaction. Every rank must execute the same + # sequence of collectives, so a rank that finds a problem does NOT raise + # where it finds it -- it records an outcome and raises only at the next + # gather, together with everyone else. Raising early would leave the + # peers blocked in a collective that their partner has already left. + self._ipc_ptrs: Optional[List[int]] = None + self._handle: Optional[int] = None + self.group = group + self.rank = dist.get_rank(group=group) + self.world_size = dist.get_world_size(group=group) + self.device = torch.device("cuda", torch.cuda.current_device()) + # Bound on first executing use; see _check_stream. + self._stream: Optional[torch.cuda.Stream] = None + self.elem_size = 0 + self.max_numel = max_numel + self.max_blocks = max_blocks + self.profile = "" + self.profile_reason = "" + # Resolved launch configurations, keyed exactly. Consulted before any + # AutoTuner call because even a pure cache lookup there takes a global + # lock, which is real overhead at this operator's scale. + self._tuned: Dict[Tuple[int, int, torch.dtype], IpcLaunchConfig] = {} + self._runner: Optional[PcieIpcAllReduceRunner] = None + self._tune_group: Optional[ProcessGroup] = None + self._tune_batches = tuple(int(b) for b in tune_batches) + self._tune_cache = tune_cache or default_cache_path(self.world_size) + self._tune_cache_exists = False + self._tuned_configs_loaded = False + self._warned_untuned = False + + # --- stage 1: local validation, encoded rather than raised ----------- + error: Optional[str] = None + if self.world_size not in _SUPPORTED_WORLD_SIZES: + error = ( + f"world size {self.world_size} unsupported; " + f"expected one of {_SUPPORTED_WORLD_SIZES}" + ) + elif dtype not in _SUPPORTED_DTYPES: + error = f"dtype {dtype} unsupported; expected one of {_SUPPORTED_DTYPES}" + else: + self.elem_size = torch.empty((), dtype=dtype).element_size() + pack_elems = 16 // self.elem_size + if max_numel <= 0 or max_numel % pack_elems != 0: + # The kernels address the scratch in 16-byte packs, so a + # capacity that is not a whole number of packs is rejected by + # the launcher on every call. Catch it here instead. + error = ( + f"max_numel must be a positive multiple of {pack_elems} " + f"for {dtype}, got {max_numel}" + ) + elif max_blocks <= 0: + error = f"max_blocks must be positive, got {max_blocks}" + + # Layout must be identical on every rank, or one of them reads a peer + # slab at the wrong offsets. Gather the config alongside the outcome so + # a single collective settles both. + local = { + "error": error, + "max_numel": max_numel, + "elem_size": self.elem_size, + "max_blocks": max_blocks, + "profile": profile, + # Buckets pick which shape a tuned entry is reused for, so ranks + # that disagree would resolve different configurations. + "tune_batches": self._tune_batches, + "tune_cache": self._tune_cache, + } + self._joint_check(local, "validating arguments") + + # --- stage 2: topology, then module + workspace size ----------------- + # Both before any allocation, so an unsupported topology or a failed + # JIT build costs nothing to unwind. + try: + decision = resolve_pcie_ipc_profile(group, requested=profile) + self.profile = decision.profile + self.profile_reason = decision.reason + module = get_pcie_ipc_comm_module() + nbytes = module.workspace_size( + self.world_size, max_numel, self.elem_size, max_blocks + ) + except Exception as e: # noqa: BLE001 - re-raised jointly below + nbytes = 0 + self._joint_check({"error": f"{type(e).__name__}: {e}"}, "preparing") + raise # unreachable: _joint_check raises on every rank + self._joint_check({"error": None}, "preparing") + + # --- stage 3: allocate and share, then bind -------------------------- + # NOTE: create_shared_buffer() runs its own all_gather_object and + # barrier internally. A failure *inside* it leaves the group in a state + # this constructor cannot repair; that is a property of the shared + # helper, not something worked around here. + self._ipc_ptrs = create_shared_buffer(nbytes, group=group) + bind_error: Optional[str] = None + try: + self._handle = module.init( + self._ipc_ptrs, self.rank, max_numel, self.elem_size, max_blocks + ) + # init() zeroes this rank's slab; no peer may push into it until + # every rank has done so. + torch.cuda.synchronize(self.device) + except Exception as e: # noqa: BLE001 - re-raised jointly below + bind_error = f"{type(e).__name__}: {e}" + + # Whether to tear down is a group decision: the cleanup itself contains + # barriers, so one rank must never enter it alone. + try: + self._joint_check({"error": bind_error}, "binding the workspace") + except Exception: + self.destroy() + raise + + # --- stage 4: tuned configurations, if any have been persisted ------- + # Loaded once, here, and never reloaded: a rank that picks up a file + # update its peers have not seen would choose a different kernel, and + # the group hangs rather than erroring. + try: + self._init_tuning() + except Exception: + self.destroy() + raise + dist.barrier(group=group) + + def _joint_check(self, local: dict, what: str) -> None: + """Gather per-rank outcomes and fail the whole group, or none of it. + + Raises the same error on every rank, so the caller can rely on all + ranks taking the same branch afterwards. + """ + gathered: List[Optional[dict]] = [None] * self.world_size + dist.all_gather_object(gathered, local, group=self.group) + entries = [g for g in gathered if g is not None] + + failed = {i: g["error"] for i, g in enumerate(entries) if g.get("error")} + if failed: + raise ValueError(f"pcie ipc workspace failed while {what}: {failed}") + + mismatched = { + key: [g[key] for g in entries] + for key in local + if key != "error" and len({repr(g[key]) for g in entries}) > 1 + } + if mismatched: + raise ValueError( + "every rank must build the workspace with identical arguments, " + f"but these differ across the group: {mismatched}" + ) + + @property + def handle(self) -> int: + if self._handle is None: + raise RuntimeError("workspace has been destroyed") + return self._handle + + def _check_stream(self) -> None: + """Bind the workspace to one stream, and reject use from another. + + The workspace carries mutable protocol state -- the epoch that selects + which half of the scratch a call stages through, and the arrival + counter that commits it. Both are advanced by the kernels themselves + and are only well defined if the calls that share this workspace are + totally ordered. Stream order gives that; two streams do not, and + concurrent calls would interleave their epoch reads and commits and + silently corrupt each other. + + Capture is exempt: `torch.cuda.graph` records on a side stream but + nothing executes, and the captured nodes form a linear chain that + replays in order. Replaying such a graph concurrently with other calls + on the same workspace is still unsupported and cannot be checked from + here. + """ + if torch.cuda.is_current_stream_capturing(): + return + current = torch.cuda.current_stream(self.device) + if self._stream is None: + self._stream = current + elif current != self._stream: + raise RuntimeError( + "this workspace is already bound to " + f"{self._stream}, but all_reduce was called on {current}. " + "One workspace serves one stream: its epoch and arrival " + "counters assume the calls sharing it are totally ordered. " + "Build a second workspace for the second stream." + ) + + def rebind_stream(self) -> None: + """Allow the next call to come from a different stream. + + The workspace rejects a second stream because it cannot tell "used + sequentially from another stream" from "used concurrently", and only + the latter is unsafe. A caller that knows the previous stream's work + has completed -- because it synchronized, or recorded and waited on an + event -- can say so here and move the binding. + + This is an assertion by the caller, not a check: calling it without + actually ordering the two streams reintroduces the corruption it exists + to prevent. + """ + self._stream = None + + def launch_config(self, inp: torch.Tensor) -> Optional[IpcLaunchConfig]: + """Seed launch configuration for ``inp``, or ``None`` if unsupported. + + The seed is a default, not a measurement -- see + :mod:`~flashinfer.comm.pcie_ipc_policy`. :meth:`tuned_launch_config` + is what returns a measured answer once :meth:`tune` has run. + + Depends only on shape, dtype and the workspace's own immutable + attributes, never on rank-local state: every rank must reach the same + answer or the collective deadlocks. + + Raises + ------ + ValueError + If ``inp`` is not on the workspace's device. This is deliberately + not reported as "unsupported": a caller checking :meth:`supports` + reads ``False`` as "use another backend", so answering ``False`` + here would turn a local bug into a silent fallback on one rank -- + and one rank taking a different branch hangs the rest. + """ + # Checked before the workspace state so the diagnosis is the same + # whether or not the workspace is still alive. + if inp.device != self.device: + raise ValueError( + f"input is on {inp.device} but the workspace was built on {self.device}" + ) + if self._handle is None: + return None + if inp.dtype not in _SUPPORTED_DTYPES: + return None + if inp.element_size() != self.elem_size: + return None + if not inp.is_contiguous() or inp.dim() == 0: + return None + numel = inp.numel() + if numel > self.max_numel: + return None + return get_pcie_ipc_launch_config( + self.world_size, numel, self.elem_size, self.max_blocks + ) + + def supports(self, inp: torch.Tensor) -> bool: + """Whether the kernels can run ``inp`` at all. + + A capability question -- dtype, contiguity, workspace capacity, and + enough payload for the reduce-scatter to give every rank a share. It + does not mean the shape has been measured on this machine; call + :meth:`tune` for that. + + Raises the same way :meth:`launch_config` does on a device mismatch -- + that is a caller bug, not an unsupported shape. + + Autotuning never changes this answer: it only picks a faster + configuration for a shape that is already supported. + """ + return self.launch_config(inp) is not None + + def _init_tuning(self) -> None: + """Build the runner and load any persisted configurations. Collective.""" + self._runner = PcieIpcAllReduceRunner(self) + path = self._tune_cache + exists = os.path.isfile(path) + # Whether the file is there has to be a group fact before anyone acts + # on it: half a group running tuned configurations and half running the + # seed is a hang, not a slowdown. + self._joint_check({"error": None, "cache": exists}, "checking the tune cache") + self._tune_cache_exists = exists + if exists: + from ..autotuner import AutoTuner + + AutoTuner.get().load_configs(path) + # Settled against the loaded keys, where the answer is known, rather + # than inferred from a miss later. + self._tuned_configs_loaded = exists and cache_covers_workspace( + self.world_size, self.profile, self.max_blocks, self.max_numel + ) + self._joint_check( + { + "error": None, + "digest": self._cache_digest(), + "covers": self._tuned_configs_loaded, + }, + "loading the tune cache", + ) + + def _warn_if_untuned(self) -> None: + """Say once that this workspace resolved to seed configurations. + + Two causes with different fixes, so two messages: a machine nobody + tuned, or a cache keyed for a different workspace (see + :func:`~flashinfer.comm.pcie_ipc_tuning.cache_covers_workspace`). + + On the cold path only, so the steady state is untouched: a serving loop + reaches this at most once per distinct shape, and the flag makes it once + per workspace. Warning here rather than in ``__init__`` keeps it tied to + actually using the kernels, not to building a workspace the caller may + never route to. + """ + if self._tuned_configs_loaded or self._warned_untuned: + return + self._warned_untuned = True + if self._tune_cache_exists: + warnings.warn( + f"PCIe IPC all-reduce loaded {self._tune_cache} but it holds no " + f"entry for this workspace ({self.world_size} ranks, " + f"max_numel={self.max_numel}, max_blocks={self.max_blocks}, " + f"profile={self.profile}); it was tuned for a different one, so " + "every shape falls back to a seed configuration. Re-tune with " + "this workspace's parameters, or build it with the ones the " + "cache was written for.", + UserWarning, + stacklevel=4, + ) + else: + warnings.warn( + "PCIe IPC all-reduce is running seed launch configurations: " + f"nothing has been tuned for {self.world_size} ranks on this " + f"machine ({self._tune_cache} does not exist). The seed picks a " + "workable kernel, not a fast one. Call workspace.tune([hidden]) " + "once per machine; the result is persisted and later processes " + "pick it up.", + UserWarning, + stacklevel=4, + ) + + def _cache_digest(self) -> str: + """Fingerprint of the tuned entries this rank will actually use. + + ``load_configs`` silently drops entries whose metadata does not match + the machine, so "we all read the same file" is not the same as "we all + hold the same table". + """ + from ..autotuner import AutoTuner + + prefix = f"('{PCIE_IPC_CUSTOM_OP}'" + tuner = AutoTuner.get() + entries = sorted( + (key, repr(value)) + for key, value in tuner._file_configs.items() + if key.startswith(prefix) + ) + return hashlib.sha256(repr(entries).encode()).hexdigest()[:16] + + def tuned_launch_config(self, inp: torch.Tensor) -> Optional[IpcLaunchConfig]: + """Launch configuration for ``inp``, measured if one has been persisted. + + Admission is asked first and is final: a shape the kernels cannot run + returns ``None`` here too, whatever the cache holds. + + Inside an ``autotune(True)`` context this runs the search; outside one + it is a lookup. Same split as the other tunable ops in this library. + """ + seed = self.launch_config(inp) + if seed is None: + return None + from ..autotuner import AutoTuner + + tuner = AutoTuner.get() + key = (inp.numel(), inp.shape[-1], inp.dtype) + # The hot cache is skipped while tuning, so a search that has more + # shapes to cover is not short-circuited by an earlier answer. + if not tuner.is_tuning_mode: + cached = self._tuned.get(key) + if cached is not None: + return cached + # Resolving is collective and reads the verdict back to the host, so it + # cannot happen inside a graph capture. Say that here: the CUDA-level + # failure is "Cannot copy between CPU and CUDA tensors during CUDA + # graph capture", which names neither this workspace nor the fix. + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"the launch configuration for shape {tuple(inp.shape)} dtype " + f"{inp.dtype} has not been resolved yet, and resolving it " + "inside a CUDA graph capture is not possible: the ranks agree " + "on it with a collective whose result is read back on the " + "host. Call workspace.prepare() with every shape you intend to " + "capture -- after tune(), which clears this cache -- or pass " + "config= explicitly at the call site." + ) + config = self._resolve_tuned(inp, seed, tuner) + self._tuned[key] = config + return config + + def _resolve_tuned( + self, inp: torch.Tensor, seed: IpcLaunchConfig, tuner + ) -> IpcLaunchConfig: + """Cold path: search or look up, then make the group agree.""" + hidden = inp.shape[-1] + batch = inp.numel() // hidden + tuning_config = pcie_ipc_tuning_config(self._tune_batches) + can_profile = tuner.is_tuning_mode and self._runner.can_profile(inp.device) + if can_profile: + _, tactic = tuner.choose_one( + PCIE_IPC_CUSTOM_OP, [self._runner], tuning_config, [inp] + ) + else: + # An enclosing autotune context may belong to another operator and + # may have replaced the global file cache. Without a matching + # distributed tune group, profiling this collective is unsafe. + # Restore the workspace's explicit tune cache and perform a lookup + # with its own bucket policy instead of retaining the seed tactic. + if tuner.is_tuning_mode: + # The runner would have said this had it been asked for + # candidates; short-circuiting before ``choose_one`` is what + # would otherwise swallow it. It is the actionable half of the + # diagnosis -- the caller is tuning, so "go tune" is not. + warn_no_tune_group(stacklevel=4) + if self._tune_cache_exists: + tuner.load_configs(self._tune_cache) + _, _, tactic, _ = tuner.search_cache( + PCIE_IPC_CUSTOM_OP, + [self._runner], + ((batch, hidden),), + tuning_config, + inputs=[inp], + ) + config = resolve_tuned_config(seed, tactic, self.world_size, self.max_blocks) + + # Unconditional, even when the cache missed and `config is seed`. The + # ranks would otherwise have to agree on whether to run this collective + # before running it, and disagreeing about that is the hang it exists + # to prevent. It costs one small reduction per distinct shape. + packed = pack_config(config) + bounds = torch.tensor([packed, -packed], dtype=torch.int64, device=self.device) + dist.all_reduce(bounds, op=dist.ReduceOp.MAX, group=self.group) + if int(bounds[0]) != -int(bounds[1]): + # Fall back rather than raise: the seed is a pure function, so it + # is agreed by construction and the group stays alive. + warnings.warn( + "ranks resolved different tuned configurations for shape " + f"{tuple(inp.shape)}; falling back to the seed configuration. " + "The tune cache is inconsistent across ranks -- delete " + f"{self._tune_cache} and re-tune.", + RuntimeWarning, + stacklevel=3, + ) + return seed + if not tuner.is_tuning_mode: + self._warn_if_untuned() + return config + + @flashinfer_api(trace=pcie_ipc_all_reduce_trace) + def all_reduce( + self, + inp: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + config: Optional[IpcLaunchConfig] = None, + enable_pdl: bool = False, + ) -> torch.Tensor: + """Out-of-place all-reduce. + + Parameters + ---------- + inp : torch.Tensor + Contiguous CUDA tensor whose byte size is a multiple of 16. + out : torch.Tensor, optional + Destination. Allocated when omitted. + config : IpcLaunchConfig, optional + Launch geometry and kernel selection. Resolved from the tune cache + or the seed when omitted; pass one explicitly only to benchmark or + to reach a kernel neither would choose. Ranks that disagree on it + hang -- see the collective contract in the class docstring. + enable_pdl : bool + Programmatic dependent launch. **Currently rejected.** The TP8 + block kernel triggers launch completion before it writes its + island ack and barrier flag, so a dependent kernel could start + while this call's protocol state is still being written. + + Returns + ------- + torch.Tensor + The reduced tensor. + + Raises + ------ + ValueError + If the kernels cannot run this shape. Check :meth:`supports` first + and fall back to another backend. + """ + if config is None: + config = self.tuned_launch_config(inp) + if config is None: + raise ValueError( + f"unsupported shape {tuple(inp.shape)} dtype {inp.dtype} " + f"at {self.world_size} ranks; check supports() first" + ) + self._check_stream() + # Raise rather than fall back: a device mismatch is a caller bug, and + # silently opting this rank out would hang every other rank. + if inp.device != self.device: + raise ValueError( + f"input is on {inp.device} but the workspace was built on {self.device}" + ) + if out is None: + out = torch.empty_like(inp) + elif out.device != self.device: + raise ValueError( + f"output is on {out.device} but the workspace was built on " + f"{self.device}" + ) + self._launch(inp, out, config, enable_pdl) + return out + + def _launch( + self, + inp: torch.Tensor, + out: torch.Tensor, + config: IpcLaunchConfig, + enable_pdl: bool = False, + ) -> None: + """Issue one collective with an explicit configuration. + + The launch without the admission, device and stream checks around it. + Callers that have already done those -- the tuner, which sweeps many + configurations over one validated pair of buffers -- use this so the + checks do not run once per candidate. + """ + get_pcie_ipc_comm_module().all_reduce( + self.handle, + inp, + out, + config.blocks, + config.threads, + int(config.variant), + enable_pdl, + ) + + def prepare( + self, + shapes: Sequence[Tuple[int, int]], + *, + dtype: torch.dtype = torch.bfloat16, + ) -> Dict[Tuple[int, int], Optional[IpcLaunchConfig]]: + """Resolve the launch configuration for each shape now. **Collective.** + + Resolution is lazy by default: the first call at a given shape looks the + configuration up and then makes the group agree on it, which costs one + small reduction whose verdict is read back on the host. That is fine in + eager mode and impossible inside a CUDA graph capture, so a shape first + used inside a capture fails to capture. + + This moves that work to a point the caller chooses. Nothing else + changes -- the same lookup, the same agreement, the same number of + collectives -- and afterwards every listed shape is served from the + in-process cache, so a capture of it touches no collective at all. + + Call it **after** :meth:`tune`, which clears that cache, and list every + shape that will be captured: shapes left out are still resolved lazily + and still cannot be captured. Serving frameworks pad the batch to the + bucket they capture, so list the padded sizes, not the real ones. + + Parameters + ---------- + shapes : sequence of (batch, hidden) + Shapes to resolve. Every rank must pass the same list in the same + order -- resolution is collective, so a rank with a different list + deadlocks the group rather than disagreeing. + dtype : torch.dtype + Which of the two supported dtypes to resolve for. The cache is + keyed by dtype, so resolve each one that will be used. + + Returns + ------- + dict + ``{(batch, hidden): config}``, with ``None`` for shapes the kernels + do not support -- those fall back to another backend at call time + and never reach a capture. + """ + shapes = [(int(batch), int(hidden)) for batch, hidden in shapes] + # Same reasoning as tune(): the loop below issues one collective per + # shape, so a rank with a different list hangs rather than disagrees. + self._joint_check( + {"error": None, "shapes": shapes, "dtype": str(dtype)}, + "preparing launch configurations", + ) + resolved: Dict[Tuple[int, int], Optional[IpcLaunchConfig]] = {} + for batch, hidden in shapes: + probe = torch.empty((batch, hidden), dtype=dtype, device=self.device) + resolved[(batch, hidden)] = self.tuned_launch_config(probe) + return resolved + + def tune( + self, + hiddens: Sequence[int], + *, + dtype: torch.dtype = torch.bfloat16, + cache: Optional[str] = None, + tune_group=None, + warmup: int = TUNE_WARMUP, + repeat: int = TUNE_REPEAT, + ) -> Dict[Tuple[int, int], IpcLaunchConfig]: + """Measure the launch configuration for every tuned shape. Collective. + + A convenience wrapper around the library's usual tuning idiom:: + + with flashinfer.autotune(True, cache=path): + for batch in batches: + ws.all_reduce(sample(batch)) + + which also works, and does the same thing. This adds what a collective + needs on top of it: a gloo subgroup for the timing reduction so every + rank picks the same kernel, longer timing runs than the library default + (the library defaults resolve too little at this scale), a check that + every rank agrees on the arguments, and a single writer for the result + file. + + Every rank must call this with identical arguments, and clocks should be + pinned first (``nvidia-smi -lgc``): boost drift is larger than the + differences being ranked. + + Parameters + ---------- + hiddens : Sequence[int] + Hidden sizes to tune -- the ones this job will actually run. There + is no default: admission does not constrain the hidden size, so + there is no finite set to enumerate, and guessing would quietly tune + a shape nobody uses. + + The **batch** dimension is not here. It comes from ``tune_batches`` + on the constructor, because the buckets have to be the same on the + tuning side and the lookup side, which makes them a property of the + workspace rather than of one call. + dtype : torch.dtype + Which of the two supported dtypes to measure. Both are 2 bytes so + the traffic is identical, but they take different conversion paths. + cache : str, optional + Where to persist results. Defaults to the workspace's + ``tune_cache``, which is also where the next process reads them. + tune_group : ProcessGroup, optional + Group used to reduce per-candidate timings so every rank picks the + same winner. Built here as a gloo subgroup when the workspace spans + the default process group; must be supplied otherwise, because + ``new_group`` is collective over the *default* group and building + one here would hang a job whose workspace is a strict subgroup. + warmup, repeat : int + Untimed and timed iterations per candidate. The library defaults + time too short a span to resolve candidates for a collective this + fast, so these default higher. + + Returns + ------- + dict + ``{(hidden, batch): config}`` for every shape that was measured, so + the caller can see what tuning actually covered and what it chose. + + Raises + ------ + ValueError + If none of ``hiddens`` yields a shape the kernels admit -- otherwise + the call is a silent no-op. + """ + from ..autotuner import ( + AutoTuner, + autotune, + get_autotune_process_group, + set_autotune_process_group, + ) + + hiddens = tuple(int(h) for h in hiddens) + path = cache or self._tune_cache + # Everything the collective profiling contract requires to match, in + # one gather. A blocklist set on one rank alone silently shortens that + # rank's candidate list, and the timing reduction then deadlocks on the + # first divergence. + self._joint_check( + { + "error": None, + "hiddens": hiddens, + "dtype": str(dtype), + "cache": path, + "warmup": warmup, + "repeat": repeat, + "tune_batches": self._tune_batches, + "blocklist": os.environ.get("FLASHINFER_TACTICS_BLOCKLIST", ""), + "digest": self._cache_digest(), + }, + "starting a tuning run", + ) + + if tune_group is None: + tune_group = self._make_tune_group() + elif dist.get_world_size(tune_group) != self.world_size: + raise ValueError( + f"tune_group spans {dist.get_world_size(tune_group)} ranks but " + f"the workspace spans {self.world_size}" + ) + + tuner = AutoTuner.get() + previous_group = get_autotune_process_group() + previous_counts = (tuner.warmup, tuner.repeat) + set_autotune_process_group(tune_group) + # The library defaults time too short a span to resolve candidates at + # this operator's scale. + tuner.warmup, tuner.repeat = warmup, repeat + covered: List[Tuple[int, int]] = [] + skipped: List[int] = [] + try: + for hidden in hiddens: + batches = [ + b + for b in tuned_batches_for( + hidden, self._tune_batches, self.max_numel + ) + if self.launch_config( + torch.empty((b, hidden), dtype=dtype, device=self.device) + ) + is not None + ] + if not batches: + # Recorded rather than skipped silently: the call would + # otherwise return cleanly having measured nothing. + skipped.append(hidden) + continue + torch.cuda.synchronize(self.device) + self.rebind_stream() + with autotune(True, tuning_buckets=tuple(batches), round_up=False): + for batch in batches: + inp = torch.randint( + 0, + 16, + (batch, hidden), + dtype=torch.int32, + device=self.device, + ).to(dtype) + tuner.choose_one( + PCIE_IPC_CUSTOM_OP, + [self._runner], + pcie_ipc_tuning_config(self._tune_batches), + [inp], + ) + covered.append((hidden, batch)) + finally: + tuner.warmup, tuner.repeat = previous_counts + # Restore rather than clear: a caller may be tuning something else + # around this. + set_autotune_process_group(previous_group) + + if skipped: + message = ( + f"tune() measured nothing for hidden {skipped} at " + f"{self.world_size} ranks: the kernels do not support those " + "shapes, and tuning does not widen what is supported." + ) + if not covered: + raise ValueError(message) + warnings.warn(message, RuntimeWarning, stacklevel=2) + + # Winners live in the in-memory cache now, so drop anything this + # workspace resolved from the seed. + self._tuned.clear() + self._tuned_configs_loaded = True + dist.barrier(group=self.group) + if self.rank == 0: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + tuner.save_configs(path) + # Nobody leaves before the file is on disk: a peer that rebuilt its + # workspace first would load a half-written table. + dist.barrier(group=self.group) + return { + (hidden, batch): self.tuned_launch_config( + torch.empty((batch, hidden), dtype=dtype, device=self.device) + ) + for hidden, batch in covered + } + + def _make_tune_group(self): + """A gloo subgroup for reducing candidate timings. + + gloo because the reduction carries one float64 and an NCCL collective + immediately after a spin-waiting IPC kernel is exactly the interference + a timing loop does not want. + """ + if self._tune_group is not None: + return self._tune_group + ranks = dist.get_process_group_ranks(self.group) + if len(ranks) != dist.get_world_size(): + raise ValueError( + "tune() cannot build its own reduction group for a workspace " + "that spans a strict subgroup: new_group() is collective over " + "the default process group, so every process would have to " + "call it. Pass tune_group= built by all ranks instead." + ) + self._tune_group = dist.new_group(ranks=ranks, backend="gloo") + return self._tune_group + + def destroy(self) -> None: + """Release the handle and the shared slab. + + Collective: every rank must call this, and the peer unmapping is + separated from the free by a barrier inside ``free_shared_buffer``. + """ + if self._handle is not None: + # all_reduce() launches asynchronously, so a collective may still be + # running or spinning on this slab. free_shared_buffer() unmaps the + # peers, and unmapping memory a live kernel is still touching is a + # use-after-free -- wait for the device before tearing anything + # down. This is the conservative choice; a stream-scoped wait would + # need the workspace to track every stream it has been used on. + torch.cuda.synchronize(self.device) + get_pcie_ipc_comm_module().dispose(self._handle) + self._handle = None + if self._ipc_ptrs is not None: + free_shared_buffer(self._ipc_ptrs, group=self.group) + self._ipc_ptrs = None + if self._tune_group is not None: + dist.destroy_process_group(self._tune_group) + self._tune_group = None + self._tuned.clear() + + def __enter__(self) -> "PcieIpcAllReduceWorkspace": + return self + + def __exit__(self, *exc_info) -> None: + self.destroy() diff --git a/flashinfer/comm/pcie_ipc_policy.py b/flashinfer/comm/pcie_ipc_policy.py new file mode 100644 index 00000000000..da181fb1f62 --- /dev/null +++ b/flashinfer/comm/pcie_ipc_policy.py @@ -0,0 +1,188 @@ +""" +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. + +Launch configurations for the PCIe IPC all-reduce. + +Two layers, with very different standing. + +**Admission** (:func:`_admits`) is a capability question: which shapes the +kernels can run at all. It is not a performance judgement, and tuning cannot +change it. + +**The seed** (:func:`_seed`) is a *default*, not a measurement. It picks the +one side of the one crossover that ports between machines -- push straight to +every peer while the payload is small, reduce-scatter/all-gather once it is +not -- and nothing finer. + +Thresholds fitted per batch on one machine do not survive the trip to another, +so only the shape of the answer lives here; the numbers come from +:meth:`~flashinfer.comm.PcieIpcAllReduceWorkspace.tune`, which measures them +where they will run. Running untuned is warned about once per workspace. +""" + +from dataclasses import dataclass, replace +from enum import IntEnum +from functools import lru_cache +from typing import Optional + + +# Block counts above this are never useful on either fabric and the workspace +# is sized for it. +MAX_BLOCKS = 128 + + +class IpcVariant(IntEnum): + """Which kernel to launch; mirrors ``fi::Variant`` in the header. + + Values cross the FFI boundary as integers, so they are append-only. + ``FLAT_STAGED`` is accepted at world size 8 only -- at 4 it would name the + same kernel as ``STAGED``, and at 2 there is no staged-vs-flat distinction. + """ + + UNSTAGED = 0 + STAGED = 1 + STAGED_RING = 2 + FLAT_STAGED = 3 + + +@dataclass(frozen=True) +class IpcLaunchConfig: + blocks: int + threads: int + variant: IpcVariant + + +# Payload above which reduce-scatter/all-gather beats pushing to every peer. +# Keyed on bytes, not tokens: the crossover trades bytes moved against barrier +# latency, and only that ratio ports between fabrics. +_SEED_STAGE_BYTES = 32 * 1024 + +# The neighbour-ordered kernel has one outbound stream per rank whatever the +# grid, so extra blocks pay only once there are bytes enough to keep the link +# busy. Capped low: with no switch-local peer, concurrent transfers collapse +# rather than add. +_SEED_RING_BYTES_PER_BLOCK = 256 * 1024 +_SEED_RING_MAX_BLOCKS = 4 + + +def _admits(world_size: int, numel: int, elem_size: int) -> bool: + """Whether the kernels can run this shape at all. + + Independent of which kernel is chosen: tuning launches every variant on + whatever shape is admitted, so a precondition that held only for the + variant the seed happens to pick would still deadlock the group under + :meth:`~flashinfer.comm.PcieIpcAllReduceWorkspace.tune`. + """ + if world_size not in (2, 4, 8): + return False + pack_elems = 16 // elem_size + # Matches the launcher's own check; the kernels address whole 16-byte packs. + if numel % pack_elems != 0: + return False + # Reduce-scatter gives each rank num_packs // world_size packs. Below one + # pack per rank that split degenerates onto a single owner: correct, but it + # leaves the other ranks idle, and a payload that small is better served by + # another backend than by an IPC collective. + return numel >= pack_elems * world_size + + +def _seed( + world_size: int, numel: int, elem_size: int, max_blocks: int +) -> IpcLaunchConfig: + """Default configuration for a shape nothing has measured yet.""" + payload = numel * elem_size + + if world_size == 2: + # Staging moves the same bytes it would have pushed, so there is no + # crossover here and no second branch to justify. + return IpcLaunchConfig(min(16, max_blocks), 128, IpcVariant.UNSTAGED) + + if payload >= _SEED_STAGE_BYTES: + # Neighbour-ordered rather than all-to-all: with no switch-local peer, + # simultaneous writes to every peer collapse, and the penalty grows with + # the payload. Picking wrong on this arm is unbounded rather than merely + # slow, which is why the threshold sits low. + blocks = max( + 1, + min( + _SEED_RING_MAX_BLOCKS, + payload // _SEED_RING_BYTES_PER_BLOCK, + max_blocks, + ), + ) + return IpcLaunchConfig(blocks, 256, IpcVariant.STAGED_RING) + + if world_size == 4: + # Staging always cuts egress at four ranks and the all-to-all form pays + # only two barriers, so the one-shot push is never the answer. One + # block, because its grid multiplies the concurrency that collapses. + return IpcLaunchConfig(1, 256, IpcVariant.STAGED) + + # Eight ranks below the crossover: the island-partitioned push has no + # barriers, which is what the staged path's six island barriers must beat. + return IpcLaunchConfig(min(16, max_blocks), 256, IpcVariant.UNSTAGED) + + +def _is_launchable(world_size: int, config: IpcLaunchConfig, max_blocks: int) -> bool: + """Reject configurations the kernels cannot accept. + + A violation here degrades to "unsupported shape" and a caller fallback, + which is far better than reaching the kernel and failing a hard check. + """ + if not 0 < config.blocks <= max_blocks: + return False + if not world_size <= config.threads <= 1024: + return False + # One configuration must name exactly one kernel, so the pairs the header + # does not dispatch are rejected rather than aliased onto a neighbour. + if world_size == 2 and config.variant not in ( + IpcVariant.UNSTAGED, + IpcVariant.STAGED, + ): + return False + if config.variant == IpcVariant.FLAT_STAGED and world_size != 8: + return False + # The block-partitioned TP8 kernel derives its chunk from blockIdx.x & 3. + # Every other kernel uses flat grid-stride loops. + if ( + world_size == 8 + and config.variant == IpcVariant.STAGED + and config.blocks % 4 != 0 + ): + return False + return True + + +@lru_cache(maxsize=None) +def get_pcie_ipc_launch_config( + world_size: int, + numel: int, + elem_size: int, + max_blocks: int = MAX_BLOCKS, +) -> Optional[IpcLaunchConfig]: + """Launch configuration for one shape, or ``None`` when unsupported. + + ``None`` means the kernels cannot run the shape, so the caller must use + another backend. It never means "untuned": an untuned shape gets the seed. + + Depends only on its arguments, so every rank in a group reaches the same + answer -- a prerequisite, since a rank that opts out while its peers opt in + deadlocks the collective. + """ + if not _admits(world_size, numel, elem_size): + return None + config = _seed(world_size, numel, elem_size, max_blocks) + config = replace(config, blocks=min(config.blocks, max_blocks)) + return config if _is_launchable(world_size, config, max_blocks) else None diff --git a/flashinfer/comm/pcie_ipc_topology.py b/flashinfer/comm/pcie_ipc_topology.py new file mode 100644 index 00000000000..1bb869afcb9 --- /dev/null +++ b/flashinfer/comm/pcie_ipc_topology.py @@ -0,0 +1,210 @@ +""" +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. +""" + +import socket +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +# Which fabric the group is on. The distinction is the interconnect, not the +# GPU: the same card behaves differently depending on whether its NUMA island +# contains a PCIe switch. It selects no kernel -- it keys the tune cache, so two +# topologies on one machine do not read each other's measurements. +PROFILE_ROOTCPLX = "rootcplx-noswitch" +PROFILE_SWITCHPAIR = "pcieswitch-pairs" +PCIE_IPC_PROFILES = (PROFILE_ROOTCPLX, PROFILE_SWITCHPAIR) + +_PROFILE_ALIASES = { + "rootcplx": PROFILE_ROOTCPLX, + "rootcplx-noswitch": PROFILE_ROOTCPLX, + "pcieswitch": PROFILE_SWITCHPAIR, + "pcieswitch-pairs": PROFILE_SWITCHPAIR, +} + + +@dataclass +class PcieIpcRankTopology: + """Per-rank probe result, exchanged across the group. + + ``peer_switch_local`` is keyed by the *peer GPU's UUID* so the decision + layer can join results across ranks regardless of each process's + ``CUDA_VISIBLE_DEVICES`` ordering, and so the probe only ever describes + GPUs this rank can actually see. + """ + + rank: int + hostname: str = "" + device_index: int = -1 + device_uuid: str = "" + peer_switch_local: Dict[str, bool] = field(default_factory=dict) + pair_errors: Dict[str, str] = field(default_factory=dict) + probe_error: Optional[str] = None + + +@dataclass(frozen=True) +class PcieIpcProfileDecision: + profile: str + reason: str + + +def probe_pcie_ipc_rank_topology( + rank: int, device: Optional[torch.device] = None +) -> PcieIpcRankTopology: + """Probe whether this rank's GPU shares a PCIe switch with any peer. + + Never raises: any failure is recorded in ``probe_error`` and the decision + layer treats an unknown topology conservatively. + + Only the GPU this rank owns is probed against the other visible GPUs, so a + job pinned to a subset of the machine describes that subset rather than the + whole host. That matters on a mixed box where one island sits behind a + switch and another does not. + """ + topo = PcieIpcRankTopology(rank=rank) + try: + topo.hostname = socket.gethostname() + parsed = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if parsed.type != "cuda": + raise ValueError(f"probe requires a CUDA device, got {parsed!r}") + device_index = ( + parsed.index if parsed.index is not None else torch.cuda.current_device() + ) + topo.device_index = device_index + + import pynvml + + pynvml.nvmlInit() + try: + + def _uuid(idx: int) -> str: + props = torch.cuda.get_device_properties(idx) + uuid = getattr(props, "uuid", None) + if uuid is None: + raise RuntimeError( + "torch.cuda.get_device_properties(...).uuid unavailable; " + "cannot establish physical GPU identity" + ) + return f"GPU-{uuid}" + + def _handle(idx: int): + return pynvml.nvmlDeviceGetHandleByUUID(_uuid(idx).encode()) + + topo.device_uuid = _uuid(device_index) + my_handle = _handle(device_index) + # NVML_TOPOLOGY_HOSTBRIDGE is the first level that leaves the switch + # fabric, so anything strictly below it means the pair talks through + # a PCIe switch without reaching the host bridge. + hostbridge = pynvml.NVML_TOPOLOGY_HOSTBRIDGE + for peer in range(torch.cuda.device_count()): + if peer == device_index: + continue + peer_uuid = _uuid(peer) + try: + level = pynvml.nvmlDeviceGetTopologyCommonAncestor( + my_handle, _handle(peer) + ) + topo.peer_switch_local[peer_uuid] = level < hostbridge + except pynvml.NVMLError as pair_err: + topo.pair_errors[peer_uuid] = str(pair_err) + finally: + pynvml.nvmlShutdown() + except Exception as e: # noqa: BLE001 - any probe failure => conservative fallback + topo.probe_error = f"{type(e).__name__}: {e}" + return topo + + +def decide_pcie_ipc_profile( + requested: Optional[str], topologies: List[PcieIpcRankTopology] +) -> PcieIpcProfileDecision: + """Pick the fabric label from the gathered probes. Pure function. + + An explicit ``requested`` profile always wins. Otherwise the group is + switch-paired only if some rank positively observed a switch-local peer; + anything unknown or unprobeable falls back to ``rootcplx-noswitch``. + Guessing wrong costs a tune cache keyed on the other fabric, so the label + that claims less is the safe default. + """ + # The intra-node constraint is checked first: CUDA IPC cannot cross hosts, + # so an explicit profile must not be able to wave it through. + hosts = {t.hostname for t in topologies if t.hostname} + if len(hosts) > 1: + raise ValueError( + f"pcie ipc all-reduce is intra-node only, but the group spans {sorted(hosts)}" + ) + + if requested is not None: + key = requested.strip().lower() + if key not in _PROFILE_ALIASES: + raise ValueError( + f"unknown pcie ipc profile {requested!r}; " + f"expected one of {sorted(_PROFILE_ALIASES)}" + ) + return PcieIpcProfileDecision(_PROFILE_ALIASES[key], "requested explicitly") + + failed = [t.rank for t in topologies if t.probe_error] + if failed: + return PcieIpcProfileDecision( + PROFILE_ROOTCPLX, f"probe failed on ranks {failed}; assuming no switch pair" + ) + + # Only pairs where BOTH endpoints belong to this group count. The probe + # walks every GPU the process can see, which for a subgroup is a superset: + # a switch-local pair outside the group says nothing about how the group's + # own ranks talk to each other. + members = {t.device_uuid for t in topologies if t.device_uuid} + for t in topologies: + for peer_uuid, switch_local in t.peer_switch_local.items(): + if switch_local and peer_uuid in members: + return PcieIpcProfileDecision( + PROFILE_SWITCHPAIR, + f"rank {t.rank} shares a PCIe switch with group member {peer_uuid}", + ) + + partial = [t.rank for t in topologies if any(u in members for u in t.pair_errors)] + if partial: + return PcieIpcProfileDecision( + PROFILE_ROOTCPLX, + f"some in-group pairs unprobeable on ranks {partial}; " + "assuming no switch pair", + ) + return PcieIpcProfileDecision(PROFILE_ROOTCPLX, "no switch-local pair observed") + + +def resolve_pcie_ipc_profile( + group: ProcessGroup, + requested: Optional[str] = None, + device: Optional[torch.device] = None, +) -> PcieIpcProfileDecision: + """Probe every rank and agree on one profile. + + Collective. Runs before any workspace allocation or JIT build so an + unsupported topology costs nothing, and gathers the per-rank probes so + every rank reaches the same decision from the same evidence. + """ + rank = dist.get_rank(group=group) + local = probe_pcie_ipc_rank_topology(rank, device=device) + gathered: List[Optional[PcieIpcRankTopology]] = [None] * dist.get_world_size( + group=group + ) + dist.all_gather_object(gathered, local, group=group) + return decide_pcie_ipc_profile(requested, [t for t in gathered if t is not None]) diff --git a/flashinfer/comm/pcie_ipc_tuning.py b/flashinfer/comm/pcie_ipc_tuning.py new file mode 100644 index 00000000000..ba4bf66b13f --- /dev/null +++ b/flashinfer/comm/pcie_ipc_tuning.py @@ -0,0 +1,530 @@ +""" +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. + +Autotuning for the PCIe IPC all-reduce. + +The seed in :mod:`~flashinfer.comm.pcie_ipc_policy` is a default, not a +measurement: one crossover, and no constants fitted to any machine. This module +measures the same choice, over the launch configurations the dispatch can +actually reach. + +Two properties of the surrounding code shape everything here: + +**The autotuner never looks at a kernel's output**, and this kernel family's +characteristic failure is wrong *and* fast. So every candidate is verified +against a reference before it is timed, and the verdict is reduced across the +group -- see :meth:`PcieIpcAllReduceRunner.get_valid_tactics`. + +**Every wait in the kernels is an unbounded spin.** Ranks that disagree on the +launch configuration, or that issue different numbers of calls, hang rather +than raise. So the candidate list is a pure function of group-identical +arguments, the verification verdict is reduced before it is used, and the +resolved configuration is checked for group agreement before it is cached. + +The policy module keeps three jobs here: admission decides which shapes are +supported at all, and the seed is both tactic ``-1`` and the fallback whenever a +tuned answer cannot be used. +""" + +import os +import warnings +from functools import lru_cache +from typing import Dict, List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist + +from ..autotuner import ( + DynamicTensorSpec, + TunableRunner, + TuningConfig, + make_bucket_mapper, +) +from .pcie_ipc_policy import ( + MAX_BLOCKS, + IpcLaunchConfig, + IpcVariant, + _is_launchable, +) + +# Baked into every persisted cache key, so renaming it silently invalidates +# every cache file rather than mis-resolving one. +PCIE_IPC_CUSTOM_OP = "flashinfer::pcie_ipc_all_reduce" + +# Bump when a variant's meaning, the scratch-region assignment, or the +# candidate encoding changes. The autotuner's own metadata records library and +# driver versions but nothing about this op, and a dev checkout does not move +# the FlashInfer version. +PCIE_IPC_TUNE_VERSION = 1 + +# Not all powers of two: the extra entries are block counts the search selected +# on real hardware, and it cannot converge on a configuration its own grid +# cannot name. +TUNE_BLOCKS: Tuple[int, ...] = (1, 2, 4, 8, 12, 16, 32, 64, 96, 128) +TUNE_THREADS: Tuple[int, ...] = (64, 128, 256, 512, 1024) + +# Batch buckets. Floor semantics, so a bucket is always a batch the tuner +# actually measured. Matches the benchmark's default sweep. +TUNE_BATCHES: Tuple[int, ...] = (1, 2, 4, 8, 16, 32, 64, 128) + +# Higher than the library defaults, which time too short a span to resolve +# candidates for a collective of this scale. +TUNE_WARMUP = 10 +TUNE_REPEAT = 50 + +# Reference tactic. The autotuner reserves -1 for "the fallback that implements +# any shape"; here that is the policy module's seed configuration. +TABLE_TACTIC = -1 + +# Inputs are drawn from [0, INIT_MAX_VALUE) so the group sum stays integral and +# exactly representable, which is what lets verification use a zero tolerance +# despite the kernels summing in a different order than NCCL. +INIT_MAX_VALUE = 16 + + +def candidate_tactics( + world_size: int, + max_blocks: int = MAX_BLOCKS, + blocks: Tuple[int, ...] = TUNE_BLOCKS, + threads: Tuple[int, ...] = TUNE_THREADS, +) -> Tuple[Tuple[int, int, int], ...]: + """Every launch configuration the dispatch can reach, as tactics. + + A pure function of group-identical arguments, so every rank derives the + same list in the same order -- which the autotuner's collective profiling + requires and cannot check. + """ + return _candidate_tactics_cached(world_size, max_blocks, blocks, threads) + + +@lru_cache(maxsize=None) +def _candidate_tactics_cached(world_size, max_blocks, blocks, threads): + out = [] + for variant in IpcVariant: + for b in blocks: + for t in threads: + if _is_launchable( + world_size, IpcLaunchConfig(b, t, variant), max_blocks + ): + out.append((int(variant), b, t)) + return tuple(out) + + +def config_to_tactic(config: IpcLaunchConfig) -> Tuple[int, int, int]: + """Encode a configuration as a tactic. + + Plain ints, because a tactic has to survive a JSON round-trip: the + autotuner writes ``[0, 32, 128]`` and reads back ``(0, 32, 128)``. + Self-describing rather than an index into :func:`candidate_tactics`, so + editing the grid cannot repoint a persisted entry at a different kernel. + """ + return (int(config.variant), int(config.blocks), int(config.threads)) + + +def tactic_to_config(tactic: Sequence[int]) -> IpcLaunchConfig: + """Decode a tactic. Raises ``ValueError`` on anything malformed.""" + if len(tactic) != 3: + raise ValueError(f"expected a 3-element tactic, got {tactic!r}") + variant, blocks, threads = (int(v) for v in tactic) + try: + return IpcLaunchConfig(blocks, threads, IpcVariant(variant)) + except ValueError as exc: + raise ValueError(f"tactic {tactic!r} names no variant: {exc}") from exc + + +def cache_covers_workspace( + world_size: int, profile: str, max_blocks: int, max_numel: int +) -> bool: + """Whether the loaded cache holds any entry written for this workspace. + + ``max_numel`` is part of the key, so a workspace sized differently from the + tuned one misses every entry at once rather than a few -- a configuration + mistake rather than an untuned shape, and no single lookup can tell those + apart, since the seed is a valid answer either way. + + dtype is not compared: one workspace serves both 2-byte dtypes and each gets + its own entries, so a match would be required for a cache that covers the + workspace perfectly well in the dtype the caller is not using. + + Scanned rather than parsed for the same reason + :meth:`PcieIpcAllReduceWorkspace._cache_digest` scans -- the key format + belongs to the autotuner. + """ + from ..autotuner import AutoTuner + + prefix = f"('{PCIE_IPC_CUSTOM_OP}'" + # cache_key_extras up to the dtype, with the closing paren traded for the + # separator that must follow it. + head = ( + PCIE_IPC_TUNE_VERSION, + int(world_size), + str(profile), + int(max_blocks), + int(max_numel), + ) + needle = repr(head)[:-1] + ", " + return any( + key.startswith(prefix) and needle in key + for key in AutoTuner.get()._file_configs + ) + + +def resolve_tuned_config( + table_config: IpcLaunchConfig, + tactic, + world_size: int, + max_blocks: int, +) -> IpcLaunchConfig: + """Turn a tactic into a configuration, falling back to the seed. + + The autotuner does not check that a cached tactic can implement the shape + it is being reused for, so a cache written against a larger ``max_blocks`` + would otherwise reach the launcher's hard checks and raise on every rank in + the middle of a collective. + """ + if tactic is None or tactic == TABLE_TACTIC: + return table_config + try: + config = tactic_to_config(tactic) + except (TypeError, ValueError): + return table_config + if not _is_launchable(world_size, config, max_blocks): + return table_config + return config + + +def small_int_initializer( + shapes: Tuple[int, ...], dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """Synthesize profiling inputs that can be compared at zero tolerance. + + The autotuner's default fills tensors with ``rand() * 10 - 5``, which no + reference can be compared against exactly. Small integers keep the group + sum exact in both supported dtypes, so verification uses ``torch.equal`` + and cannot mistake a reduction-order difference for a protocol bug. Zero is + in the range on purpose: the sentinel kernels rewrite real zeros in the + payload, and that path should be exercised. + """ + return torch.randint( + 0, INIT_MAX_VALUE, shapes, device=device, dtype=torch.int32 + ).to(dtype) + + +@lru_cache(maxsize=None) +def pcie_ipc_tuning_config(batches: Tuple[int, ...] = TUNE_BATCHES) -> TuningConfig: + """Tuning configuration for one bucket set. + + Cached so that the serving-side cache lookup and the tuning-side search + share one object: the bucket mapper has to be identity-stable or the + autotuner's profile lookup degenerates. + + Only the batch dimension is dynamic. Hidden stays static so it lands + verbatim in the cache key -- the configuration follows the payload in bytes, + and a bucketed hidden would silently reuse another payload's answer. + """ + return TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=(0,), + dim_idx=(0,), + gen_tuning_buckets=batches, + map_to_tuning_buckets=make_bucket_mapper(batches, round_map=False), + ), + ), + tensor_initializers=((0, small_int_initializer),), + # Capture is required, not preferred. Without it the profiler issues + # each iteration separately, the host cannot keep up with a collective + # this short, and the span it times is dominated by launch gaps -- it + # would rank host overhead rather than kernels. + use_cold_l2_cache=False, + use_cuda_graph=True, + ) + + +def default_cache_path(world_size: int) -> str: + """Where tuned configurations are persisted. + + World size is in the filename as well as the cache key so that a TP4 and a + TP8 job on the same host never contend for one file. + """ + import pathlib + + override = os.getenv("FLASHINFER_AUTOTUNE_DIR") + if override: + base = pathlib.Path(override) + else: + from ..jit.env import FLASHINFER_WORKSPACE_DIR + + base = FLASHINFER_WORKSPACE_DIR / "autotune" + return str(base / f"pcie_ipc_all_reduce_ws{world_size}.json") + + +def cache_key_extras( + world_size: int, + profile: str, + max_blocks: int, + max_numel: int, + dtype: torch.dtype, +) -> Tuple: + """Everything the autotuner's own cache key leaves out. + + That key is only the bucketed input shapes, so without these a TP4 and a + TP8 entry at the same shape would collide, a configuration tuned on one + fabric would be reused on the other, and a cache written for one workspace + size would be applied to another. ``max_numel`` matters because the epoch + double buffer places its halves ``world_size * max_numel`` apart, so the + best block count genuinely depends on it. + + Every field is a workspace immutable or the input dtype, which is what the + autotuner requires: the tuple must come out the same for the caller's real + tensors and for the ones it synthesizes. + """ + return ( + PCIE_IPC_TUNE_VERSION, + int(world_size), + str(profile), + int(max_blocks), + int(max_numel), + str(dtype), + ) + + +def pack_config(config: IpcLaunchConfig) -> int: + """Pack a configuration into one integer for a cross-rank comparison.""" + return ( + (int(config.variant) << 32) | (int(config.blocks) << 16) | int(config.threads) + ) + + +def reduce_verdict(wrong: "torch.Tensor", group) -> "torch.Tensor": + """Make every rank agree on which candidates computed the wrong answer. + + ``MAX`` over a per-candidate "was wrong" flag, which is the same decision + as ``MIN`` over "was right": one rank seeing a mismatch condemns the + candidate everywhere. A rank-local verdict would let ranks profile + different candidate sets, and the autotuner's timing reduction then + deadlocks on the first divergence. + + Corruption is not necessarily uniform across ranks: the cross-island race + this protocol can produce leaves some of them clean, so a rank-local verdict + can miss it entirely. + + Factored out so a test can assert the operator without a GPU. + """ + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + return wrong + + +def tuned_batches_for( + hidden: int, batches: Tuple[int, ...], max_numel: int +) -> Tuple[int, ...]: + """Drop buckets that would exceed the workspace at this hidden size.""" + return tuple(b for b in batches if b * hidden <= max_numel) + + +def warn_no_tune_group(stacklevel: int = 2) -> None: + """Say why a tuning session left this collective untuned. + + Raised from two places that reach the same dead end -- the runner, when the + autotuner does ask it for candidates, and the workspace, when it declines + to ask at all. The generic "nothing is tuned" advice does not fit here: the + caller *is* tuning, so telling them to tune is a dead end. What they are + missing is the reduction group, and that is what this names. + """ + warnings.warn( + "PCIe IPC all-reduce skipped autotuning: no matching " + "autotune process group is installed on every rank. Call " + "PcieIpcAllReduceWorkspace.tune(), or install one with " + "set_autotune_process_group() before entering autotune().", + RuntimeWarning, + stacklevel=stacklevel, + ) + + +class PcieIpcAllReduceRunner(TunableRunner): + """Adapts the all-reduce to the autotuner, and screens candidates first. + + One instance per workspace, built once and kept: the autotuner puts + ``hash(runner)`` in its in-memory cache key, so a fresh instance per call + would miss every entry and re-tune. + """ + + def __init__(self, workspace) -> None: + # A weak-ish coupling on purpose: the runner needs the raw launch and + # the group, not the public API, whose admission checks would run once + # per candidate and whose tracing decorator would recurse. + self._ws = workspace + # Named to end in _cache so the base __hash__ would skip it even if the + # override below is ever removed. + self._buf_cache: Dict[Tuple[Tuple[int, ...], torch.dtype], torch.Tensor] = {} + + def __hash__(self) -> int: + # Everything that changes what this runner does, and nothing that + # changes per call. The base implementation hashes __dict__ values and + # would fold in the workspace object's identity, which differs between + # processes and would defeat the persisted cache. + ws = self._ws + return hash( + ( + type(self).__name__, + PCIE_IPC_TUNE_VERSION, + ws.world_size, + ws.profile, + ws.max_blocks, + ws.max_numel, + ) + ) + + def get_cache_key_extras(self, inputs) -> Tuple: + ws = self._ws + return cache_key_extras( + ws.world_size, ws.profile, ws.max_blocks, ws.max_numel, inputs[0].dtype + ) + + def _output_for(self, inp: torch.Tensor) -> torch.Tensor: + key = (tuple(inp.shape), inp.dtype) + out = self._buf_cache.get(key) + if out is None: + out = torch.empty_like(inp) + self._buf_cache[key] = out + return out + + def _table_config(self, inp: torch.Tensor) -> Optional[IpcLaunchConfig]: + return self._ws.launch_config(inp) + + def can_profile(self, device) -> bool: + """Whether a real search is safe, as a group decision. + + Reduced rather than read locally because the answer decides how many + times each rank enters the profiler. Ranks that search different numbers + of candidates do not disagree, they deadlock. + """ + from ..autotuner import get_autotune_process_group + + group = get_autotune_process_group() + ok = group is not None and dist.get_world_size(group) == self._ws.world_size + flag = torch.tensor([1 if ok else 0], dtype=torch.int32, device=device) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self._ws.group) + return bool(flag.item()) + + def get_valid_tactics(self, inputs, profile) -> List: + """Candidates that computed the right answer, in a group-agreed order. + + This is the gate the autotuner does not have. It selects by ``argmin`` + on wall time and never inspects an output, while this kernel family's + characteristic failure -- a sentinel poll returning stale data rather + than waiting -- is wrong *and* fast. Screening here rather than during + profiling keeps the verdict's collective out of the timed window, and + costs one launch per candidate on a cache miss only. + + Cardinality is the hazard. Every rank must issue exactly these launches + in exactly this order; an early return between the first launch and the + verdict reduction leaves peers spinning inside a kernel this rank never + issued, with no timeout. Hence: barrier first, every buffer allocated + before the loop, and a loop body that does not allocate, synchronise + with the host, or branch. + """ + inp = inputs[0] + ws = self._ws + table_config = self._table_config(inp) + if table_config is None: + # The autotuner is being asked about a shape the kernels cannot run + # at all. Nothing to choose between; the caller falls back. + return [TABLE_TACTIC] + + if not self.can_profile(inp.device): + # Tuning mode is process-global, so this op can be swept by a + # caller that only meant to tune its GEMMs. Without a reduction + # over the candidate timings the ranks would argmin independently + # and pick different kernels, which this protocol does not survive. + # Offering only the seed degrades that into a no-op. + warn_no_tune_group(stacklevel=3) + return [TABLE_TACTIC] + + tactics = candidate_tactics(ws.world_size, ws.max_blocks) + configs = [table_config] + [tactic_to_config(t) for t in tactics] + + ref = inp.clone() + dist.all_reduce(ref, group=ws.group) + out = self._output_for(inp) + wrong = torch.zeros(len(configs), dtype=torch.int32, device=inp.device) + + dist.barrier(group=ws.group) + for i, config in enumerate(configs): + # A kernel that leaves part of the payload unwritten would + # otherwise show the previous candidate's correct result. + out.fill_(float("nan")) + ws._launch(inp, out, config) + wrong[i] = torch.ne(out, ref).any() + reduce_verdict(wrong, ws.group) + + verdict = wrong.tolist() + if verdict[0]: + # The seed computing the wrong answer is not something to route + # around: it is what every untuned shape and every cache miss falls + # back to. The verdict is group-wide, so + # every rank raises together and the group unwinds cleanly. + raise RuntimeError( + "the seed configuration for shape " + f"{tuple(inp.shape)} ({table_config}) does not match a " + "reference all-reduce; refusing to tune on top of it" + ) + survivors = zip(tactics, verdict[1:], strict=True) + return [TABLE_TACTIC] + [t for t, bad in survivors if not bad] + + def forward( + self, inputs, tactic=TABLE_TACTIC, do_preparation: bool = False, **kwargs + ): + inp = inputs[0] + out = self._output_for(inp) + if do_preparation: + # Buffer now allocated; launching here would make the call counts + # depend on whether the autotuner decided to prepare. + return out + table_config = self._table_config(inp) + if table_config is None: + raise RuntimeError( + f"shape {tuple(inp.shape)} is not one the kernels support; " + "the tuner must not have been asked about it" + ) + config = resolve_tuned_config( + table_config, tactic, self._ws.world_size, self._ws.max_blocks + ) + self._ws._launch(inp, out, config) + return out + + +__all__ = [ + "PCIE_IPC_CUSTOM_OP", + "PcieIpcAllReduceRunner", + "PCIE_IPC_TUNE_VERSION", + "TABLE_TACTIC", + "TUNE_BATCHES", + "TUNE_BLOCKS", + "TUNE_REPEAT", + "TUNE_THREADS", + "TUNE_WARMUP", + "cache_key_extras", + "candidate_tactics", + "config_to_tactic", + "default_cache_path", + "pack_config", + "pcie_ipc_tuning_config", + "reduce_verdict", + "resolve_tuned_config", + "small_int_initializer", + "tactic_to_config", + "tuned_batches_for", +] diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index b2ed55bfd28..aa28cedf43c 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -84,6 +84,7 @@ from .comm import gen_trtllm_mnnvl_comm_module as gen_trtllm_mnnvl_comm_module from .comm import gen_trtllm_comm_module as gen_trtllm_comm_module from .comm import gen_vllm_comm_module as gen_vllm_comm_module +from .comm import gen_pcie_ipc_comm_module as gen_pcie_ipc_comm_module from .comm import gen_moe_alltoall_module as gen_moe_alltoall_module from .comm import gen_dcp_alltoall_module as gen_dcp_alltoall_module from .dsv3_optimizations import ( diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 3facecf44ad..4aeb49134db 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -211,6 +211,66 @@ def gen_vllm_comm_module() -> JitSpec: ) +def gen_pcie_ipc_comm_debug_module( + stall_ns: int, + stall_island: int, + no_block_epoch: int = 0, + per_block_epoch: int = 0, + no_barrier_entry_sync: int = 0, +) -> JitSpec: + """Build the kernels with one protocol mechanism disabled. Not for shipping. + + Used by the opt-in negative-control tests and by the benchmark's + ``--protocol-ab`` mode. Every switch produces an INCORRECT build. + + ``stall_ns``/``stall_island`` + Stall one island between the owner-pair rendezvous and the cross read. + The only way to open the TP8 cross-island window -- a host-side delay + cannot, because the pair barrier releases both islands together. + + ``no_block_epoch`` + Pin the scratch to half 0, i.e. no double buffer. Matches what the two + staged TP8 kernels used to do; for the others it is a protocol that + never shipped. + + ``per_block_epoch`` + Pick the half by per-block parity instead of per call. Matches what + TP2, TP4 and the TP8 pack kernel used to do. One change in grid size + then desynchronises the block ranges permanently. + + ``no_barrier_entry_sync`` + Drop the leading CTA barrier from the three signalling helpers. Applies + to every kernel that takes a barrier. + """ + return gen_jit_spec( + f"pcie_ipc_comm_dbg{stall_ns}_{stall_island}_{no_block_epoch}" + f"_{per_block_epoch}_{no_barrier_entry_sync}", + [ + jit_env.FLASHINFER_CSRC_DIR / "pcie_ipc_all_reduce.cu", + ], + extra_cuda_cflags=[ + f"-DFLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS={stall_ns}", + f"-DFLASHINFER_PCIE_IPC_DEBUG_STALL_ISLAND={stall_island}", + f"-DFLASHINFER_PCIE_IPC_DEBUG_NO_BLOCK_EPOCH={no_block_epoch}", + f"-DFLASHINFER_PCIE_IPC_DEBUG_PER_BLOCK_EPOCH={per_block_epoch}", + f"-DFLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC={no_barrier_entry_sync}", + ], + ) + + +def gen_pcie_ipc_comm_module() -> JitSpec: + # No architecture restriction: the kernels use only plain PTX loads/stores + # and CUDA IPC, both of which predate every architecture flashinfer builds + # for. The target is a PCIe machine without NVLink, which is orthogonal to + # the SM version. + return gen_jit_spec( + "pcie_ipc_comm", + [ + jit_env.FLASHINFER_CSRC_DIR / "pcie_ipc_all_reduce.cu", + ], + ) + + def gen_ulysses_a2a_module() -> JitSpec: return gen_jit_spec( "ulysses_a2a", diff --git a/flashinfer/trace/templates/comm.py b/flashinfer/trace/templates/comm.py index 4a4ea6b7752..7bfaf537f26 100644 --- a/flashinfer/trace/templates/comm.py +++ b/flashinfer/trace/templates/comm.py @@ -282,3 +282,81 @@ def _decode_cp_a2a_alltoall_init( reference=_decode_cp_a2a_alltoall_reference, init=_decode_cp_a2a_alltoall_init, ) + + +# ── PCIe IPC all-reduce (intra-node, no NVLink) ────────────────────────────── + + +@torch.no_grad() +def _pcie_ipc_all_reduce_reference( + inp: torch.Tensor, + *, + out: torch.Tensor = None, + config=None, + enable_pdl: bool = False, +) -> torch.Tensor: + """Single-rank reference: an all-reduce over one rank is the identity. + + Same modelling choice as ``allreduce_fusion`` above -- the trace runs in a + single process, so the cross-rank reduction cannot be exercised here. + Multi-rank correctness is covered by + ``tests/comm/test_pcie_ipc_all_reduce.py``, which compares against NCCL at + zero tolerance. + """ + return inp.clone() if out is None else out.copy_(inp) + + +def _pcie_ipc_all_reduce_init( + *, + num_tokens: int, + hidden_dim: int = 6144, + device: str = "cuda", + seed: int = 0, +): + """Build this rank's input for ``PcieIpcAllReduceWorkspace.all_reduce``. + + The workspace itself is an opaque multi-rank IPC handle bound to ``self`` + and is not built here; see ``tests/comm/test_pcie_ipc_all_reduce.py``. + """ + generator = torch.Generator(device=device).manual_seed(seed) + return { + "inp": torch.randn( + num_tokens, + hidden_dim, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + } + + +pcie_ipc_all_reduce_trace = TraceTemplate( + op_type="comm", + name_prefix="pcie_ipc_all_reduce", + description=( + "Custom all-reduce for intra-node PCIe machines without NVLink. The " + "launch configuration follows the payload in bytes -- a seed default " + "until the workspace is tuned on the machine it runs on -- so the " + "traced axes are the ones that select it." + ), + axes={ + "num_tokens": Var(description="Token count along dim 0."), + "hidden_dim": Const(abbrev="h"), + }, + inputs={ + "inp": Tensor( + ["num_tokens", "hidden_dim"], + description="Pre-reduction token activations (this rank's shard).", + ), + }, + outputs={ + "output": Tensor( + ["num_tokens", "hidden_dim"], + dtype_from="inp", + description="Reduced activations.", + ), + }, + tags=["status:verified", "stage:comm"], + reference=_pcie_ipc_all_reduce_reference, + init=_pcie_ipc_all_reduce_init, +) diff --git a/include/flashinfer/comm/pcie_ipc_all_reduce.cuh b/include/flashinfer/comm/pcie_ipc_all_reduce.cuh new file mode 100644 index 00000000000..0f989d83c6b --- /dev/null +++ b/include/flashinfer/comm/pcie_ipc_all_reduce.cuh @@ -0,0 +1,2247 @@ +/* + * 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. + */ +#ifndef FLASHINFER_COMM_PCIE_IPC_ALL_REDUCE_CUH_ +#define FLASHINFER_COMM_PCIE_IPC_ALL_REDUCE_CUH_ + +// Custom all-reduce for intra-node PCIe machines without NVLink. +// +// Every peer transfer on such a machine crosses the CPU root complex, where +// all-to-all writes collapse to a fraction of what the same kernel achieves +// when each rank writes to a single destination. The kernels here therefore +// stage their pushes so that at any instant each rank has exactly one outbound +// and one inbound stream, and the 8-rank path keeps a 4+4 island decomposition +// so the scarce cross-socket links carry the minimum traffic. See the PR +// description for the bandwidth measurements this is derived from; they are a +// property of the machine, not of the code. +// +// All state lives in a caller-owned workspace shared over CUDA IPC; see +// compute_workspace_layout() for the byte layout and make_peer_views() for the +// per-region pointers. The caller owns the allocation because tearing it down +// needs a collective barrier between "every rank unmaps its peers" and "every +// rank frees its own slab", which a destructor cannot express. + +#include +#include +#include + +#include +#include +#include + +namespace flashinfer { +namespace comm { +namespace pcie_ipc { + +constexpr int kMaxWorldSize = 8; +constexpr int kSignalPhases = 8; + +// Which kernel all_reduce() launches, together with world_size. Values are +// part of the FFI signature, so they are explicit and append-only. +// +// kFlatStaged is accepted at world_size 8 only: at 4 it would name the same +// kernel as kStaged, and at 2 there is no staged-vs-flat distinction. +enum class Variant : int { + kUnstaged = 0, // push to every peer at once + kStaged = 1, // staged pushes; island-decomposed at world_size 8 + kStagedRing = 2, // staged pushes in neighbour order; world_size 4 and 8 + kFlatStaged = 3, // staged pushes without the island decomposition +}; + +constexpr int kVariantCount = 4; + +// Which staging area a kernel uses. The kernels come in two protocol families +// and a region may hold only one of them. +// +// Sentinel kernels poll for +0.0 meaning "not yet written", sanitise real zeros +// out of the payload, and store +0.0 back once a poll succeeds. Barrier kernels +// are content-blind: they publish raw payload and leave it there. Nothing else +// sweeps the workspace -- the host zeroes it once at init and never again. +// +// So a sentinel kernel landing on a barrier kernel's leftovers reads stale +// payload, and its all-gather poll, which watches a single slot, exits on it +// immediately: wrong output, not a hang. The epoch double buffer does not +// substitute for this -- it guarantees the other half is quiescent, not clean. +// +// At world_size 8 that puts the two topology kernels in kBlock and both +// sentinel kernels in kPack. +enum class ScratchRegion : int { kBlock = 0, kPack = 1 }; + +template +struct alignas(sizeof(T) * N) Vec { + T data[N]; +}; + +template +struct PackTraits { + static constexpr int kPackElems = 16 / sizeof(T); + using Pack = Vec; +}; + +template +__device__ __forceinline__ void pdl_grid_sync_const() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (Enabled) { + cudaGridDependencySynchronize(); + } +#endif +} + +template +__device__ __forceinline__ void pdl_grid_release_const() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (Enabled) { + __syncthreads(); + __threadfence(); + if (threadIdx.x == 0) { + cudaTriggerProgrammaticLaunchCompletion(); + } + } +#endif +} + +__device__ __forceinline__ void store_release_i32(int32_t* addr, int32_t value) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(value), "l"(addr)); +#else + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(value), "l"(addr)); +#endif +} + +__device__ __forceinline__ int32_t load_acquire_i32(int32_t* addr) { + int32_t value; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(value) : "l"(addr)); +#else + asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" : "=r"(value) : "l"(addr)); +#endif + return value; +} + +// Has the peer reached generation `expected`? +// +// Generations are a free-running counter, so a plain `observed < expected` +// breaks the first time it wraps: the slot still holds the old maximum, the +// new generation has wrapped to the minimum, and the comparison lets the +// barrier through before the peer has arrived. Switching to unsigned does not +// fix it either -- it just moves the break to UINT_MAX -> 0. +// +// Compare on the circle instead: reinterpret the difference as a signed +// distance modulo 2^32. `observed - expected >= 0` is then true exactly when +// the peer is at or past `expected`, for any pair within 2^31 generations of +// each other -- which is always, since a rank advances one generation per +// call and its peers are at most one call behind. +// constexpr so the boundary behaviour can be pinned at compile time, below. +// Those static_asserts are the whole defence against this being "simplified" +// back to a plain `<`: that version is correct for two billion calls and then +// releases every barrier a generation early, which no test anyone would +// actually run is going to catch. +__host__ __device__ __forceinline__ constexpr bool generation_reached(int32_t observed, + int32_t expected) { + return static_cast(static_cast(observed) - static_cast(expected)) >= + 0; +} + +static_assert(generation_reached(5, 5), "a peer at the expected generation has arrived"); +static_assert(generation_reached(6, 5), "a peer past the expected generation has arrived"); +static_assert(!generation_reached(4, 5), "a peer one generation behind has not arrived"); +// The wrap that motivated this function. Signed `observed < expected` reads +// INT32_MAX < INT32_MIN as false and lets the barrier through. +static_assert(generation_reached(INT32_MIN, INT32_MAX), + "the generation after INT32_MAX has arrived"); +static_assert(!generation_reached(INT32_MAX, INT32_MIN), + "the generation before the wrap has not arrived"); +// Unsigned `<` fixes the pair above but breaks this one, at UINT32_MAX -> 0. +// Only the modular comparison gets both. +static_assert(generation_reached(0, -1), "0 is one generation past -1"); +static_assert(!generation_reached(-1, 0), "-1 is one generation before 0"); + +__device__ __forceinline__ void store_volatile_i32(int32_t* addr, int32_t value) { + asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(value), "l"(addr)); +} + +__device__ __forceinline__ int32_t load_volatile_i32(int32_t* addr) { + int32_t value; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(value) : "l"(addr)); + return value; +} + +__device__ __forceinline__ float to_float(float x) { return x; } +__device__ __forceinline__ float to_float(half x) { return __half2float(x); } +__device__ __forceinline__ float to_float(nv_bfloat16 x) { return __bfloat162float(x); } + +template +__device__ __forceinline__ T from_float(float x); + +template <> +__device__ __forceinline__ float from_float(float x) { + return x; +} + +template <> +__device__ __forceinline__ half from_float(float x) { + return __float2half(x); +} + +template <> +__device__ __forceinline__ nv_bfloat16 from_float(float x) { + return __float2bfloat16(x); +} + +__device__ __forceinline__ uint32_t add_half2_u32(uint32_t a, uint32_t b) { + auto ah = *reinterpret_cast(&a); + auto bh = *reinterpret_cast(&b); + half2 out = __hadd2(ah, bh); + return *reinterpret_cast(&out); +} + +__device__ __forceinline__ uint32_t add_bfloat162_u32(uint32_t a, uint32_t b) { + auto ah = *reinterpret_cast<__nv_bfloat162*>(&a); + auto bh = *reinterpret_cast<__nv_bfloat162*>(&b); + __nv_bfloat162 out = __hadd2(ah, bh); + return *reinterpret_cast(&out); +} + +template +__device__ __forceinline__ uint4 packed_add_u4(uint4 a, uint4 b) { + if constexpr (std::is_same_v) { + a.x = add_half2_u32(a.x, b.x); + a.y = add_half2_u32(a.y, b.y); + a.z = add_half2_u32(a.z, b.z); + a.w = add_half2_u32(a.w, b.w); + } else { + static_assert(std::is_same_v); + a.x = add_bfloat162_u32(a.x, b.x); + a.y = add_bfloat162_u32(a.y, b.y); + a.z = add_bfloat162_u32(a.z, b.z); + a.w = add_bfloat162_u32(a.w, b.w); + } + return a; +} + +template +__device__ __forceinline__ float2 lane_to_float2(uint32_t lane); + +template <> +__device__ __forceinline__ float2 lane_to_float2(uint32_t lane) { + auto value = *reinterpret_cast(&lane); + return __half22float2(value); +} + +template <> +__device__ __forceinline__ float2 lane_to_float2(uint32_t lane) { + auto value = *reinterpret_cast<__nv_bfloat162*>(&lane); + return __bfloat1622float2(value); +} + +template +__device__ __forceinline__ uint32_t float2_to_lane(float2 value); + +template <> +__device__ __forceinline__ uint32_t float2_to_lane(float2 value) { + half2 out = __float22half2_rn(value); + return *reinterpret_cast(&out); +} + +template <> +__device__ __forceinline__ uint32_t float2_to_lane(float2 value) { + __nv_bfloat162 out = __float22bfloat162_rn(value); + return *reinterpret_cast(&out); +} + +template +__device__ __forceinline__ uint4 reduce_u4_fp32(uint4 const (&values)[WorldSize]) { + float2 acc0 = {0.0f, 0.0f}; + float2 acc1 = {0.0f, 0.0f}; + float2 acc2 = {0.0f, 0.0f}; + float2 acc3 = {0.0f, 0.0f}; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + float2 v0 = lane_to_float2(values[peer].x); + float2 v1 = lane_to_float2(values[peer].y); + float2 v2 = lane_to_float2(values[peer].z); + float2 v3 = lane_to_float2(values[peer].w); + acc0.x += v0.x; + acc0.y += v0.y; + acc1.x += v1.x; + acc1.y += v1.y; + acc2.x += v2.x; + acc2.y += v2.y; + acc3.x += v3.x; + acc3.y += v3.y; + } + uint4 out; + out.x = float2_to_lane(acc0); + out.y = float2_to_lane(acc1); + out.z = float2_to_lane(acc2); + out.w = float2_to_lane(acc3); + return out; +} + +template +struct ZeroBits; + +template <> +struct ZeroBits { + using Raw = uint16_t; + static constexpr Raw kPos = 0x0000u; + static constexpr Raw kNeg = 0x8000u; +}; + +template <> +struct ZeroBits { + using Raw = uint16_t; + static constexpr Raw kPos = 0x0000u; + static constexpr Raw kNeg = 0x8000u; +}; + +template <> +struct ZeroBits { + using Raw = uint32_t; + static constexpr Raw kPos = 0x00000000u; + static constexpr Raw kNeg = 0x80000000u; +}; + +template +__device__ __forceinline__ void clear_pos_zero(T& value) { + using Bits = ZeroBits; + using Raw = typename Bits::Raw; + Raw* raw = reinterpret_cast(&value); + if (*raw == Bits::kPos) { + *raw = Bits::kNeg; + } +} + +template +__device__ __forceinline__ bool is_pos_zero(T value) { + using Bits = ZeroBits; + using Raw = typename Bits::Raw; + Raw raw = *reinterpret_cast(&value); + return raw == Bits::kPos; +} + +template +__device__ __forceinline__ T pos_zero() { + using Bits = ZeroBits; + using Raw = typename Bits::Raw; + Raw raw = Bits::kPos; + return *reinterpret_cast(&raw); +} + +template +__device__ __forceinline__ typename PackTraits::Pack load_pack_volatile( + typename PackTraits::Pack const* base, int idx) { + uint4 raw; + auto const* addr = reinterpret_cast(base + idx); + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w) + : "l"(addr)); + return *reinterpret_cast::Pack*>(&raw); +} + +template +__device__ __forceinline__ void store_pack_volatile(typename PackTraits::Pack* base, int idx, + typename PackTraits::Pack value) { + uint4 raw = *reinterpret_cast(&value); + auto* addr = reinterpret_cast(base + idx); + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(raw.x), "r"(raw.y), + "r"(raw.z), "r"(raw.w), "l"(addr)); +} + +template +__device__ __forceinline__ void clear_pos_zero_pack(typename PackTraits::Pack& pack) { +#pragma unroll + for (int i = 0; i < PackTraits::kPackElems; ++i) { + clear_pos_zero(pack.data[i]); + } +} + +template +__device__ __forceinline__ bool has_pos_zero_pack(typename PackTraits::Pack const& pack) { + bool has_zero = false; +#pragma unroll + for (int i = 0; i < PackTraits::kPackElems; ++i) { + has_zero |= is_pos_zero(pack.data[i]); + } + return has_zero; +} + +template +__device__ __forceinline__ typename PackTraits::Pack zero_pack() { + typename PackTraits::Pack pack; +#pragma unroll + for (int i = 0; i < PackTraits::kPackElems; ++i) { + pack.data[i] = pos_zero(); + } + return pack; +} + +__device__ __forceinline__ uint32_t clear_pos_zero_u16x2(uint32_t raw) { + uint32_t lo = raw & 0xffffu; + uint32_t hi = raw & 0xffff0000u; + if (lo == 0u) { + lo = 0x8000u; + } + if (hi == 0u) { + hi = 0x80000000u; + } + return hi | lo; +} + +__device__ __forceinline__ bool has_pos_zero_u16x2(uint32_t raw) { + return (raw & 0xffffu) == 0u || (raw & 0xffff0000u) == 0u; +} + +__device__ __forceinline__ uint4 clear_pos_zero_u4_16(uint4 value) { + value.x = clear_pos_zero_u16x2(value.x); + value.y = clear_pos_zero_u16x2(value.y); + value.z = clear_pos_zero_u16x2(value.z); + value.w = clear_pos_zero_u16x2(value.w); + return value; +} + +__device__ __forceinline__ bool has_pos_zero_u4_16(uint4 value) { + return has_pos_zero_u16x2(value.x) || has_pos_zero_u16x2(value.y) || + has_pos_zero_u16x2(value.z) || has_pos_zero_u16x2(value.w); +} + +__device__ __forceinline__ uint4 load_u4_volatile(uint4 const* base, int idx) { + uint4 value; + auto const* addr = base + idx; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(value.x), "=r"(value.y), "=r"(value.z), "=r"(value.w) + : "l"(addr)); + return value; +} + +__device__ __forceinline__ void store_u4_volatile(uint4* base, int idx, uint4 value) { + auto* addr = base + idx; + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(value.x), "r"(value.y), + "r"(value.z), "r"(value.w), "l"(addr)); +} + +__device__ __forceinline__ int phase_offset(int phase, int block, int peer, int max_blocks, + int world_size) { + // Epoch slots occupy [0, max_blocks). Barrier slots start after that + // dedicated prefix so phase 0 can never corrupt an epoch. + return max_blocks + phase * max_blocks * world_size + block * world_size + peer; +} + +__device__ __forceinline__ int flag_offset(int block, int max_blocks, int world_size) { + return max_blocks + kSignalPhases * max_blocks * world_size + block; +} + +// Call-level double-buffer state: the epoch at [0] and the arrival counter at +// [1], placed past the flag region so no existing offset moves. +// +// ONE PAIR PER SCRATCH REGION. The hard constraint runs one way only: +// +// **Kernels writing the same region MUST share a counter.** TP4 alternates +// ipc_rsag_push and ipc_rsag_ring with the payload and both write the block +// region at the same addresses, so private counters would let a ring call on +// half 0 be followed immediately by a push call that also reads 0 -- back to +// back, with nothing in between to drain the first. +// +// Kernels in *different* regions need not be: any intervening collective +// already drains the previous one. They are kept apart anyway, because binding +// the state on the same host line that picks the region makes the two +// impossible to get out of step. +// +// Rank-local, and indexed by nothing, so every CTA of a launch agrees on the +// half regardless of gridDim. Per-block parity cannot: it counts how many times +// *that block* has run, so one change in gridDim desynchronises the block ranges +// permanently and a block picks a half another block is still using. +// +// Consistency across ranks comes from the same SPMD argument that already +// backs the per-block barrier flags: every rank runs the same sequence of +// collectives, so every rank is on the same call parity. +__host__ __device__ __forceinline__ int scratch_state_offset(int max_blocks, int world_size, + ScratchRegion region) { + return max_blocks + kSignalPhases * max_blocks * world_size + max_blocks + + 2 * static_cast(region); +} + +// Debug only: pin every call to half 0, i.e. the pre-double-buffer behaviour. +// Kept because the cross-island race is invisible without a way to build the +// broken protocol on demand -- it is what proves a repro actually has power, +// and it isolates the cost of the double buffer in a single benchmark session. +// Never define this in a shipping build. +// Debug only: remove the leading CTA barrier from the three signalling helpers. +// The resulting build is incorrect -- a warp can announce "my stage is done" +// while its siblings are still writing. Never define this in a shipping build. +#ifndef FLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC +#define FLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC 0 +#endif + +#ifndef FLASHINFER_PCIE_IPC_DEBUG_NO_BLOCK_EPOCH +#define FLASHINFER_PCIE_IPC_DEBUG_NO_BLOCK_EPOCH 0 +#endif + +// Debug only: restore the per-block epoch parity these kernels used before the +// call-level counter replaced it. This is the negative control for the +// grid-change regression -- without it, a passing test cannot distinguish "the +// fix works" from "the sequence never opened the window". Distinct from +// NO_BLOCK_EPOCH above, which removes double buffering entirely; this one keeps +// two halves and only breaks the agreement about which half a call is on. +// Never define this in a shipping build. +#ifndef FLASHINFER_PCIE_IPC_DEBUG_PER_BLOCK_EPOCH +#define FLASHINFER_PCIE_IPC_DEBUG_PER_BLOCK_EPOCH 0 +#endif + +// Read this call's epoch and advance it, both at kernel entry. +// +// Advancing at entry rather than at exit is what keeps this affordable. The +// flip only has to follow every block's *read*, not every block's work: the +// state is rank-local (views.self_signal), so its only reader is the next +// kernel on this stream, and that cannot start until this launch has fully +// retired. The last block to arrive therefore knows every peer block has +// already read, and can flip immediately. +// +// Committing at exit instead would put a tail __syncthreads() plus an L2 atomic +// round trip on the block-retirement critical path. It would also be fragile: +// an early return added after the arrival would freeze the counter below +// gridDim.x - 1 and silently pin the epoch. +// +// Electing the last arrival deliberately does not depend on block scheduling +// order or on the whole grid being resident. "Block 0 flips" would need every +// block launched before block 0 reaches this point, which CUDA does not +// promise, and a grid-wide spin barrier deadlocks once gridDim exceeds +// occupancy. +// +// At gridDim.x == 1 there is nothing to elect, so the atomic is skipped; the +// end state is identical (counter 0, epoch flipped). +// Exit half of the PER_BLOCK_EPOCH debug build; compiles to nothing otherwise. +// The flip sits at the exit because that is where the implementation it rebuilds +// put it, and entry-versus-exit changes the cost. +__device__ __forceinline__ void debug_commit_per_block_epoch(int32_t* per_block_slot, int epoch) { +#if FLASHINFER_PCIE_IPC_DEBUG_PER_BLOCK_EPOCH + // Bare -- no __syncthreads(), matching the implementation this rebuilds. + if (threadIdx.x == 0) { + store_volatile_i32(per_block_slot, epoch ^ 1); + } +#else + (void)per_block_slot; + (void)epoch; +#endif +} + +__device__ __forceinline__ int advance_scratch_epoch(int32_t* state, int32_t* per_block_slot) { +#if FLASHINFER_PCIE_IPC_DEBUG_NO_BLOCK_EPOCH + (void)state; + (void)per_block_slot; + return 0; +#elif FLASHINFER_PCIE_IPC_DEBUG_PER_BLOCK_EPOCH + // Read only; the flip is issued at kernel exit by + // debug_commit_per_block_epoch(). + (void)state; + return load_volatile_i32(per_block_slot) & 1; +#else + (void)per_block_slot; + const int epoch = load_volatile_i32(state) & 1; + // Every thread must have read before this block announces its arrival. + __syncthreads(); + if (threadIdx.x == 0) { + // Device scope, not system scope. This state is rank-local -- its only + // reader is this rank's next kernel on this stream -- and that reader wants + // nothing but the value itself, so there is nothing for a release to order. + // st.release.sys would flush this thread's writes system-wide over PCIe on + // the entry critical path. + if (gridDim.x == 1) { + // Sole CTA: trivially the last arrival, and the counter is already 0. + store_volatile_i32(state, epoch ^ 1); + } else if (atomicAdd(state + 1, 1) == static_cast(gridDim.x) - 1) { + store_volatile_i32(state + 1, 0); + store_volatile_i32(state, epoch ^ 1); + } + } + return epoch; +#endif +} + +__device__ __forceinline__ void block_barrier(uint64_t const* signal_ptrs, int rank, int world_size, + int max_blocks, int phase, int flag) { + // Publishing a signal means "this CTA's writes for the previous stage are + // done", so every thread must have finished them before the signalling + // threads announce it. The call sites use __threadfence_system(), which + // orders only the *calling* thread's accesses -- it does not wait for the + // rest of the CTA, and cannot help at all where the previous stage was a load + // loop. Only a CTA barrier establishes that. +#if !FLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC + __syncthreads(); +#endif + int block = blockIdx.x; + int32_t* self = reinterpret_cast(signal_ptrs[rank]); + if (threadIdx.x < world_size) { + int peer = threadIdx.x; + int32_t* peer_signal = reinterpret_cast(signal_ptrs[peer]); + store_release_i32(peer_signal + phase_offset(phase, block, rank, max_blocks, world_size), flag); + int32_t* self_slot = self + phase_offset(phase, block, peer, max_blocks, world_size); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + __syncthreads(); +} + +__device__ __forceinline__ void block_barrier_mask(uint64_t const* signal_ptrs, int rank, + int world_size, int max_blocks, int phase, + int flag, uint32_t participant_mask) { + // Entry barrier: see block_barrier above. +#if !FLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC + __syncthreads(); +#endif + if ((participant_mask & (1u << rank)) == 0u) { + __syncthreads(); + return; + } + int block = blockIdx.x; + int32_t* self = reinterpret_cast(signal_ptrs[rank]); + if (threadIdx.x < world_size) { + int peer = threadIdx.x; + if ((participant_mask & (1u << peer)) != 0u) { + int32_t* peer_signal = reinterpret_cast(signal_ptrs[peer]); + store_release_i32(peer_signal + phase_offset(phase, block, rank, max_blocks, world_size), + flag); + int32_t* self_slot = self + phase_offset(phase, block, peer, max_blocks, world_size); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + } + __syncthreads(); +} + +__device__ __forceinline__ void island_owner_gather(uint64_t const* signal_ptrs, int rank, int base, + int owner, int max_blocks, int phase, + int flag) { + // Entry barrier: see block_barrier above. +#if !FLASHINFER_PCIE_IPC_DEBUG_NO_BARRIER_ENTRY_SYNC + __syncthreads(); +#endif + int block = blockIdx.x; + int32_t* owner_signal = reinterpret_cast(signal_ptrs[owner]); + if (threadIdx.x == 0) { + store_release_i32(owner_signal + phase_offset(phase, block, rank, max_blocks, 8), flag); + } + if (rank == owner && threadIdx.x < 4) { + int peer = base + threadIdx.x; + int32_t* self_slot = owner_signal + phase_offset(phase, block, peer, max_blocks, 8); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + __syncthreads(); +} + +__device__ __forceinline__ void owner_pair_barrier(uint64_t const* signal_ptrs, int rank, int owner, + int cross_owner, int max_blocks, int phase, + int flag) { + __syncthreads(); + if (rank == owner && threadIdx.x == 0) { + int block = blockIdx.x; + int32_t* cross_signal = reinterpret_cast(signal_ptrs[cross_owner]); + store_release_i32(cross_signal + phase_offset(phase, block, rank, max_blocks, 8), flag); + int32_t* self_signal = reinterpret_cast(signal_ptrs[rank]); + int32_t* self_slot = self_signal + phase_offset(phase, block, cross_owner, max_blocks, 8); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + __syncthreads(); +} + +__device__ __forceinline__ void island_owner_ready(uint64_t const* signal_ptrs, int rank, int base, + int owner, int max_blocks, int phase, int flag) { + __syncthreads(); + int block = blockIdx.x; + if (rank == owner) { + if (threadIdx.x < 4) { + int peer = base + threadIdx.x; + int32_t* peer_signal = reinterpret_cast(signal_ptrs[peer]); + store_release_i32(peer_signal + phase_offset(phase, block, owner, max_blocks, 8), flag); + } + } else if (threadIdx.x == 0) { + int32_t* self_signal = reinterpret_cast(signal_ptrs[rank]); + int32_t* self_slot = self_signal + phase_offset(phase, block, owner, max_blocks, 8); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + __syncthreads(); +} + +__device__ __forceinline__ void island_owner_ack(uint64_t const* signal_ptrs, int rank, int base, + int owner, int max_blocks, int phase, int flag) { + __syncthreads(); + int block = blockIdx.x; + int32_t* owner_signal = reinterpret_cast(signal_ptrs[owner]); + if (rank != owner && threadIdx.x == 0) { + store_release_i32(owner_signal + phase_offset(phase, block, rank, max_blocks, 8), flag); + } + if (rank == owner && threadIdx.x < 4) { + int peer = base + threadIdx.x; + if (peer != owner) { + int32_t* self_slot = owner_signal + phase_offset(phase, block, peer, max_blocks, 8); + while (!generation_reached(load_acquire_i32(self_slot), flag)) { + } + } + } + __syncthreads(); +} + +template +__device__ __forceinline__ typename PackTraits::Pack add_pack(typename PackTraits::Pack a, + typename PackTraits::Pack b) { + using Traits = PackTraits; + using Pack = typename Traits::Pack; + if constexpr (std::is_same_v || std::is_same_v) { + uint4 av = *reinterpret_cast(&a); + uint4 bv = *reinterpret_cast(&b); + uint4 out = packed_add_u4(av, bv); + return *reinterpret_cast(&out); + } + + Pack out; +#pragma unroll + for (int i = 0; i < Traits::kPackElems; ++i) { + out.data[i] = from_float(to_float(a.data[i]) + to_float(b.data[i])); + } + return out; +} + +template +__device__ __forceinline__ typename PackTraits::Pack reduce_loaded_packs( + typename PackTraits::Pack const (&values)[WorldSize]) { + using Pack = typename PackTraits::Pack; + if constexpr (std::is_same_v || std::is_same_v) { + uint4 acc = *reinterpret_cast(&values[0]); +#pragma unroll + for (int peer = 1; peer < WorldSize; ++peer) { + uint4 next = *reinterpret_cast(&values[peer]); + acc = packed_add_u4(acc, next); + } + return *reinterpret_cast(&acc); + } else { + Pack acc = values[0]; +#pragma unroll + for (int peer = 1; peer < WorldSize; ++peer) { + acc = add_pack(acc, values[peer]); + } + return acc; + } +} + +// Debug-only hook for reproducing the cross-island scratch race. +// +// The hazard needs the SLOW island to still be reading the cross slot while +// the FAST island's next call overwrites it. Delaying a whole kernel launch +// from the host cannot produce that: the pair barrier releases both islands +// together, after which the reader reaches its cross read almost immediately +// while the writer still has several phases to go. The stall has to be here, +// between the pair rendezvous and the cross read. +// +// Enabled only when FLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS is defined, and +// only on the island selected by ..._STALL_ISLAND. Never define these in a +// shipping build. +#ifndef FLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS +#define FLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS 0 +#endif +#ifndef FLASHINFER_PCIE_IPC_DEBUG_STALL_ISLAND +#define FLASHINFER_PCIE_IPC_DEBUG_STALL_ISLAND 0 +#endif + +__device__ __forceinline__ void debug_cross_read_stall(int rank) { +#if FLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS > 0 + const int island = rank < 4 ? 0 : 1; + if (island == FLASHINFER_PCIE_IPC_DEBUG_STALL_ISLAND) { + __nanosleep(FLASHINFER_PCIE_IPC_DEBUG_CROSS_STALL_NS); + } + __syncthreads(); +#else + (void)rank; +#endif +} + +template +struct PushOneshotParamData { + uint64_t tmp_ptrs[kMaxWorldSize]; + uint64_t signal_ptrs[kMaxWorldSize]; + T const* input; + T* output; + int32_t* epoch_slots; + // Call-level double-buffer state for this launch's scratch region, bound + // host-side so the state and the region cannot disagree. See + // scratch_state_offset(). + int32_t* scratch_state; + int num_packs; + int rank_stride_packs; + int epoch_stride_packs; + int rank; + int max_blocks; +}; + +template +struct IpcTp2RemotePushData { + uint64_t tmp_ptrs[2]; + T const* input; + T* output; + int32_t* epoch_slots; + // Call-level double-buffer state for this launch's scratch region, bound + // host-side so the state and the region cannot disagree. See + // scratch_state_offset(). + int32_t* scratch_state; + int num_packs; + // Half-size of the epoch double buffer, in packs. Derived from max_numel, + // NOT from this call's num_packs: the two epoch halves must sit at fixed + // addresses. If they moved with the payload, a rank that finished a large + // collective and flipped its epoch would start writing a small one inside + // the region a lagging peer is still draining -- corrupting it, or having + // the peer's reset wipe the just-published data so the poll never ends. + // Every other v2 kernel already derives its stage offset this way. + int rank_stride_packs; + int rank; +}; + +template +__global__ __launch_bounds__(1024, 1) void ipc_tp2_remote_push_kernel( + const IpcTp2RemotePushData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + pdl_grid_sync_const(); + + int peer = params.rank ^ 1; + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + int epoch = advance_scratch_epoch(params.scratch_state, epoch_slot); + int stage_offset = epoch * 2 * params.rank_stride_packs; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + int peer_write_base = stage_offset + params.rank * params.num_packs; + int local_poll_base = stage_offset + peer * params.num_packs; + uint4 reset = {0u, 0u, 0u, 0u}; + if constexpr (Stream) { + for (int idx = tid; idx < params.num_packs; idx += stride) { + uint4 local_value = input[idx]; + uint4 publish_value = clear_pos_zero_u4_16(local_value); + store_u4_volatile(peer_buffer, peer_write_base + idx, publish_value); + uint4 peer_value; + while (true) { + peer_value = load_u4_volatile(local_buffer, local_poll_base + idx); + if (!has_pos_zero_u4_16(peer_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = packed_add_u4(local_value, peer_value); + store_u4_volatile(local_buffer, local_poll_base + idx, reset); + } + } else { + for (int idx = tid; idx < params.num_packs; idx += stride) { + uint4 value = clear_pos_zero_u4_16(input[idx]); + store_u4_volatile(peer_buffer, peer_write_base + idx, value); + } + for (int idx = tid; idx < params.num_packs; idx += stride) { + uint4 peer_value; + while (true) { + peer_value = load_u4_volatile(local_buffer, local_poll_base + idx); + if (!has_pos_zero_u4_16(peer_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = packed_add_u4(input[idx], peer_value); + store_u4_volatile(local_buffer, local_poll_base + idx, reset); + } + } + } else { + Pack const* input = reinterpret_cast(params.input); + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + int peer_write_base = stage_offset + params.rank * params.num_packs; + int local_poll_base = stage_offset + peer * params.num_packs; + Pack reset = zero_pack(); + if constexpr (Stream) { + for (int idx = tid; idx < params.num_packs; idx += stride) { + Pack local_value = input[idx]; + Pack publish_value = local_value; + clear_pos_zero_pack(publish_value); + store_pack_volatile(peer_buffer, peer_write_base + idx, publish_value); + Pack peer_value; + while (true) { + peer_value = load_pack_volatile(local_buffer, local_poll_base + idx); + if (!has_pos_zero_pack(peer_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = add_pack(local_value, peer_value); + store_pack_volatile(local_buffer, local_poll_base + idx, reset); + } + } else { + for (int idx = tid; idx < params.num_packs; idx += stride) { + Pack value = input[idx]; + clear_pos_zero_pack(value); + store_pack_volatile(peer_buffer, peer_write_base + idx, value); + } + for (int idx = tid; idx < params.num_packs; idx += stride) { + Pack peer_value; + while (true) { + peer_value = load_pack_volatile(local_buffer, local_poll_base + idx); + if (!has_pos_zero_pack(peer_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = add_pack(input[idx], peer_value); + store_pack_volatile(local_buffer, local_poll_base + idx, reset); + } + } + } + + debug_commit_per_block_epoch(epoch_slot, epoch); + pdl_grid_release_const(); +} + +// Owner of a pack under the reduce-scatter split. Must agree with the explicit +// chunk ranges the same kernels walk, which give the remainder to the last rank +// -- so `part == 0` (fewer packs than ranks) gives it the whole payload too. +// Writing to one owner and polling another spins forever: no timeout here. +template +__device__ __forceinline__ int rsag_owner_for_pack(int idx, int part) { + int owner = part > 0 ? idx / part : WorldSize - 1; + return owner < WorldSize ? owner : WorldSize - 1; +} + +// Staged (neighbour-ordered) RS/AG push. +// +// ipc_rsag_push_param_kernel below has every rank writing to all WorldSize-1 +// peers at the same time. Where every peer transfer crosses the CPU root +// complex that pattern runs far below what the same kernel-issued writes reach +// when a rank has a single outbound destination -- the collective is limited by +// the shape of the traffic, not the amount. +// +// This variant keeps the algorithm and the total bytes identical and only +// reorders the pushes: each push phase is split into WorldSize-1 passes, and in +// pass p every rank writes solely to peer (rank + 1 + p) % WorldSize. That map +// is a permutation, so during a pass every GPU has exactly one outbound stream +// and every GPU is the target of exactly one -- the pattern the fabric +// sustains. A barrier between passes keeps the ranks in the same pass, since +// without it they drift and the passes overlap back into all-to-all. +// +// The cost is 2*(WorldSize-1) barriers per collective, which is why the policy +// only selects this variant once the payload is large enough to pay for them. +// Keep the established phase numbering for this kernel. Epoch slots and +// barrier phases have separate ranges in the signal region (see +// phase_offset), so phase 1 is no longer needed for alias avoidance. +constexpr int kRingRsPhase0 = 1; + +template +__global__ __launch_bounds__(1024, 1) void ipc_rsag_ring_push_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + // RS uses phases [1, WorldSize-1], AG uses [WorldSize, 2*WorldSize-2]. + static_assert(2 * WorldSize - 2 < kSignalPhases, + "ring push needs 2*(WorldSize-1) barrier phases"); + pdl_grid_sync_const(); + + int32_t* self_signal = reinterpret_cast(params.signal_ptrs[params.rank]); + int flag = static_cast( + static_cast( + load_acquire_i32(self_signal + flag_offset(blockIdx.x, params.max_blocks, WorldSize))) + + 1u); + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + int epoch = advance_scratch_epoch(params.scratch_state, epoch_slot); + int stage_offset = epoch * params.epoch_stride_packs; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + int part = params.num_packs / WorldSize; + int my_start = params.rank * part; + int my_end = (params.rank == WorldSize - 1) ? params.num_packs : my_start + part; + int my_slot = stage_offset + params.rank * params.rank_stride_packs; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + uint4 reset = {0u, 0u, 0u, 0u}; + + // Reduce-scatter, own chunk: stays on this GPU, so it costs no fabric time + // and does not need a pass of its own. + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_u4_volatile(local_buffer, my_slot + idx, clear_pos_zero_u4_16(input[idx])); + } + + // Reduce-scatter, staged: one destination per pass. + for (int p = 0; p < WorldSize - 1; ++p) { + int target = (params.rank + 1 + p) % WorldSize; + int t_start = target * part; + int t_end = (target == WorldSize - 1) ? params.num_packs : t_start + part; + auto* target_buffer = reinterpret_cast(params.tmp_ptrs[target]); + for (int idx = t_start + tid; idx < t_end; idx += stride) { + store_u4_volatile(target_buffer, my_slot + idx, clear_pos_zero_u4_16(input[idx])); + } + block_barrier(params.signal_ptrs, params.rank, WorldSize, params.max_blocks, + kRingRsPhase0 + p, flag); + } + + // Owner reduce: every contribution is now in this rank's own buffer, so + // this phase touches local memory only. The reduced value is stashed back + // into this rank's own slot for the all-gather passes to re-read. + for (int idx = my_start + tid; idx < my_end; idx += stride) { + uint4 values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_u4_volatile(local_buffer, offset); + waiting |= has_pos_zero_u4_16(values[peer]); + } + if (!waiting) { + break; + } + } + uint4 acc = values[0]; +#pragma unroll + for (int peer = 1; peer < WorldSize; ++peer) { + acc = packed_add_u4(acc, values[peer]); + } + reinterpret_cast(params.output)[idx] = acc; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + if (peer != params.rank) { + store_u4_volatile(local_buffer, stage_offset + peer * params.rank_stride_packs + idx, + reset); + } + } + store_u4_volatile(local_buffer, my_slot + idx, clear_pos_zero_u4_16(acc)); + } + + // All-gather, staged: same permutation schedule as the reduce-scatter. + for (int p = 0; p < WorldSize - 1; ++p) { + int target = (params.rank + 1 + p) % WorldSize; + auto* target_buffer = reinterpret_cast(params.tmp_ptrs[target]); + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_u4_volatile(target_buffer, my_slot + idx, + load_u4_volatile(local_buffer, my_slot + idx)); + } + block_barrier(params.signal_ptrs, params.rank, WorldSize, params.max_blocks, WorldSize + p, + flag); + } + + // Consume the chunks owned by others, then clear this rank's own slot so + // the sentinel state is clean for the epoch that reuses it. + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + if (owner == params.rank) { + continue; + } + int offset = stage_offset + owner * params.rank_stride_packs + idx; + uint4 value; + while (true) { + value = load_u4_volatile(local_buffer, offset); + if (!has_pos_zero_u4_16(value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = value; + store_u4_volatile(local_buffer, offset, reset); + } + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_u4_volatile(local_buffer, my_slot + idx, reset); + } + } else { + Pack const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + Pack reset = zero_pack(); + + for (int idx = my_start + tid; idx < my_end; idx += stride) { + Pack value = input[idx]; + clear_pos_zero_pack(value); + store_pack_volatile(local_buffer, my_slot + idx, value); + } + for (int p = 0; p < WorldSize - 1; ++p) { + int target = (params.rank + 1 + p) % WorldSize; + int t_start = target * part; + int t_end = (target == WorldSize - 1) ? params.num_packs : t_start + part; + auto* target_buffer = reinterpret_cast(params.tmp_ptrs[target]); + for (int idx = t_start + tid; idx < t_end; idx += stride) { + Pack value = input[idx]; + clear_pos_zero_pack(value); + store_pack_volatile(target_buffer, my_slot + idx, value); + } + block_barrier(params.signal_ptrs, params.rank, WorldSize, params.max_blocks, + kRingRsPhase0 + p, flag); + } + for (int idx = my_start + tid; idx < my_end; idx += stride) { + Pack values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_pack_volatile(local_buffer, offset); + waiting |= has_pos_zero_pack(values[peer]); + } + if (!waiting) { + break; + } + } + Pack acc = reduce_loaded_packs(values); + reinterpret_cast(params.output)[idx] = acc; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + if (peer != params.rank) { + store_pack_volatile(local_buffer, stage_offset + peer * params.rank_stride_packs + idx, + reset); + } + } + Pack publish = acc; + clear_pos_zero_pack(publish); + store_pack_volatile(local_buffer, my_slot + idx, publish); + } + for (int p = 0; p < WorldSize - 1; ++p) { + int target = (params.rank + 1 + p) % WorldSize; + auto* target_buffer = reinterpret_cast(params.tmp_ptrs[target]); + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_pack_volatile(target_buffer, my_slot + idx, + load_pack_volatile(local_buffer, my_slot + idx)); + } + block_barrier(params.signal_ptrs, params.rank, WorldSize, params.max_blocks, WorldSize + p, + flag); + } + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + if (owner == params.rank) { + continue; + } + int offset = stage_offset + owner * params.rank_stride_packs + idx; + Pack value; + while (true) { + value = load_pack_volatile(local_buffer, offset); + if (!has_pos_zero_pack(value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = value; + store_pack_volatile(local_buffer, offset, reset); + } + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_pack_volatile(local_buffer, my_slot + idx, reset); + } + } + + if (threadIdx.x == 0) { + store_release_i32(self_signal + flag_offset(blockIdx.x, params.max_blocks, WorldSize), flag); + } + debug_commit_per_block_epoch(epoch_slot, epoch); + pdl_grid_release_const(); +} + +template +__global__ __launch_bounds__(1024, 1) void ipc_rsag_push_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + pdl_grid_sync_const(); + + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + int epoch = advance_scratch_epoch(params.scratch_state, epoch_slot); + int stage_offset = epoch * params.epoch_stride_packs; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + int part = params.num_packs / WorldSize; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + uint4 reset = {0u, 0u, 0u, 0u}; + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + int offset = stage_offset + params.rank * params.rank_stride_packs + idx; + uint4 value = clear_pos_zero_u4_16(input[idx]); + store_u4_volatile(owner_buffer, offset, value); + } + + int start = params.rank * part; + int end = (params.rank == WorldSize - 1) ? params.num_packs : start + part; + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + for (int idx = start + tid; idx < end; idx += stride) { + uint4 values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_u4_volatile(local_buffer, offset); + waiting |= has_pos_zero_u4_16(values[peer]); + } + if (!waiting) { + break; + } + } + uint4 acc = values[0]; +#pragma unroll + for (int peer = 1; peer < WorldSize; ++peer) { + acc = packed_add_u4(acc, values[peer]); + } + reinterpret_cast(params.output)[idx] = acc; + uint4 publish = clear_pos_zero_u4_16(acc); +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + store_u4_volatile(local_buffer, offset, reset); + if (peer != params.rank) { + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(peer_buffer, final_offset, publish); + } + } + } + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + if (owner == params.rank) { + continue; + } + int offset = stage_offset + owner * params.rank_stride_packs + idx; + uint4 value; + while (true) { + value = load_u4_volatile(local_buffer, offset); + if (!has_pos_zero_u4_16(value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = value; + store_u4_volatile(local_buffer, offset, reset); + } + } else { + Pack const* input = reinterpret_cast(params.input); + Pack reset = zero_pack(); + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + int offset = stage_offset + params.rank * params.rank_stride_packs + idx; + Pack value = input[idx]; + clear_pos_zero_pack(value); + store_pack_volatile(owner_buffer, offset, value); + } + + int start = params.rank * part; + int end = (params.rank == WorldSize - 1) ? params.num_packs : start + part; + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + for (int idx = start + tid; idx < end; idx += stride) { + Pack values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_pack_volatile(local_buffer, offset); + waiting |= has_pos_zero_pack(values[peer]); + } + if (!waiting) { + break; + } + } + Pack acc = reduce_loaded_packs(values); + reinterpret_cast(params.output)[idx] = acc; + Pack publish = acc; + clear_pos_zero_pack(publish); +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int offset = stage_offset + peer * params.rank_stride_packs + idx; + store_pack_volatile(local_buffer, offset, reset); + if (peer != params.rank) { + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(peer_buffer, final_offset, publish); + } + } + } + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int owner = rsag_owner_for_pack(idx, part); + if (owner == params.rank) { + continue; + } + int offset = stage_offset + owner * params.rank_stride_packs + idx; + Pack value; + while (true) { + value = load_pack_volatile(local_buffer, offset); + if (!has_pos_zero_pack(value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = value; + store_pack_volatile(local_buffer, offset, reset); + } + } + debug_commit_per_block_epoch(epoch_slot, epoch); +} + +template +__global__ __launch_bounds__(1024, 1) void ipc_topo_rsag8_push_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + pdl_grid_sync_const(); + + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + int epoch = advance_scratch_epoch(params.scratch_state, epoch_slot); + int stage_offset = epoch * params.epoch_stride_packs; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + int part = params.num_packs / 4; + int base = params.rank < 4 ? 0 : 4; + int cross_base = base ^ 4; + int local_rank = params.rank - base; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + uint4 reset = {0u, 0u, 0u, 0u}; + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int chunk = rsag_owner_for_pack<4>(idx, part); + int owner = base + chunk; + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + int input_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + uint4 local_value = input[idx]; + uint4 publish_input = clear_pos_zero_u4_16(local_value); + store_u4_volatile(owner_buffer, input_offset, publish_input); + + if (local_rank == chunk) { + uint4 values[4]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer_local] = load_u4_volatile(local_buffer, offset); + waiting |= has_pos_zero_u4_16(values[peer_local]); + } + if (!waiting) { + break; + } + } + uint4 local_sum = values[0]; +#pragma unroll + for (int peer_local = 1; peer_local < 4; ++peer_local) { + local_sum = packed_add_u4(local_sum, values[peer_local]); + } +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + int offset = stage_offset + peer * params.rank_stride_packs + idx; + store_u4_volatile(local_buffer, offset, reset); + } + + int cross_owner = cross_base + chunk; + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + int cross_write = stage_offset + params.rank * params.rank_stride_packs + idx; + uint4 publish_sum = clear_pos_zero_u4_16(local_sum); + store_u4_volatile(cross_buffer, cross_write, publish_sum); + + int cross_read = stage_offset + cross_owner * params.rank_stride_packs + idx; + uint4 cross_sum; + while (true) { + cross_sum = load_u4_volatile(local_buffer, cross_read); + if (!has_pos_zero_u4_16(cross_sum)) { + break; + } + } + uint4 final_value = packed_add_u4(local_sum, cross_sum); + reinterpret_cast(params.output)[idx] = final_value; + store_u4_volatile(local_buffer, cross_read, reset); + + uint4 publish_final = clear_pos_zero_u4_16(final_value); +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + if (peer == params.rank) { + continue; + } + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(peer_buffer, final_offset, publish_final); + } + } else { + int final_offset = stage_offset + owner * params.rank_stride_packs + idx; + uint4 final_value; + while (true) { + final_value = load_u4_volatile(local_buffer, final_offset); + if (!has_pos_zero_u4_16(final_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = final_value; + store_u4_volatile(local_buffer, final_offset, reset); + } + } + } else { + Pack const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + Pack reset = zero_pack(); + + for (int idx = tid; idx < params.num_packs; idx += stride) { + int chunk = rsag_owner_for_pack<4>(idx, part); + int owner = base + chunk; + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + int input_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + Pack local_value = input[idx]; + Pack publish_input = local_value; + clear_pos_zero_pack(publish_input); + store_pack_volatile(owner_buffer, input_offset, publish_input); + + if (local_rank == chunk) { + Pack values[4]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + int offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer_local] = load_pack_volatile(local_buffer, offset); + waiting |= has_pos_zero_pack(values[peer_local]); + } + if (!waiting) { + break; + } + } + Pack local_sum = values[0]; +#pragma unroll + for (int peer_local = 1; peer_local < 4; ++peer_local) { + local_sum = add_pack(local_sum, values[peer_local]); + } +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + int offset = stage_offset + peer * params.rank_stride_packs + idx; + store_pack_volatile(local_buffer, offset, reset); + } + + int cross_owner = cross_base + chunk; + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + int cross_write = stage_offset + params.rank * params.rank_stride_packs + idx; + Pack publish_sum = local_sum; + clear_pos_zero_pack(publish_sum); + store_pack_volatile(cross_buffer, cross_write, publish_sum); + + int cross_read = stage_offset + cross_owner * params.rank_stride_packs + idx; + Pack cross_sum; + while (true) { + cross_sum = load_pack_volatile(local_buffer, cross_read); + if (!has_pos_zero_pack(cross_sum)) { + break; + } + } + Pack final_value = add_pack(local_sum, cross_sum); + reinterpret_cast(params.output)[idx] = final_value; + store_pack_volatile(local_buffer, cross_read, reset); + + Pack publish_final = final_value; + clear_pos_zero_pack(publish_final); +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + if (peer == params.rank) { + continue; + } + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(peer_buffer, final_offset, publish_final); + } + } else { + int final_offset = stage_offset + owner * params.rank_stride_packs + idx; + Pack final_value; + while (true) { + final_value = load_pack_volatile(local_buffer, final_offset); + if (!has_pos_zero_pack(final_value)) { + break; + } + } + reinterpret_cast(params.output)[idx] = final_value; + store_pack_volatile(local_buffer, final_offset, reset); + } + } + } + + debug_commit_per_block_epoch(epoch_slot, epoch); + pdl_grid_release_const(); +} + +// TP8 staged topology RS/AG push. +// +// ipc_topo_rsag8_block_param_kernel below partitions blocks by blockIdx.x & 3, +// so its four block groups push to four different island owners at the same +// instant: intra-island traffic is all-to-all, which is the expensive pattern +// on this fabric. +// +// This variant keeps the topology decomposition exactly as it is (island reduce +// -> owner-pair exchange across SYS -> island gather), because a topology-blind +// ring is slower at every block count. It only re-times the two +// intra-island phases: each is split into three passes, and in pass p a rank +// talks solely to island peer (local + 1 + p) % 4, which is a permutation, so +// each GPU has one outbound stream at a time. The cross-island exchange is +// already one-to-one and is left alone. +// +// Because chunks are now visited in time rather than assigned to block groups, +// blocks no longer have to be a multiple of four and a flat grid-stride loop +// covers each chunk. +// +// Costs six extra island barriers, so the policy only selects this above a +// payload threshold. +template +__global__ __launch_bounds__(1024, 1) void ipc_topo_rsag8_ring_push_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + static_assert(kSignalPhases >= 8, "topology ring push needs eight barrier phases"); + pdl_grid_sync_const(); + + int32_t* self_signal = reinterpret_cast(params.signal_ptrs[params.rank]); + // Unsigned arithmetic for the bump: signed overflow is UB, and this counter + // is meant to wrap. generation_reached() reads it back on the circle. + int flag = + static_cast(static_cast(load_acquire_i32( + self_signal + flag_offset(blockIdx.x, params.max_blocks, 8))) + + 1u); + + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = gridDim.x * blockDim.x; + const int part = params.num_packs >> 2; + const int base = params.rank < 4 ? 0 : 4; + const int local = params.rank & 3; + const uint32_t island_mask = params.rank < 4 ? 0x0fu : 0xf0u; + const int cross_owner = params.rank ^ 4; + // Call-level double buffer. The cross-island payload this rank publishes + // into its paired owner's slab has no read-complete edge coming back: phase + // 3 proves the write landed, but nothing stops the paired owner's *next* + // call from overwriting it while this one is still reading. Alternating + // halves supplies the missing distance, and two halves are exactly enough -- + // owner_pair_barrier is a two-sided rendezvous, so this rank cannot leave + // call k until its partner has entered call k, hence cannot reach call k+2 + // (the next use of this half) until the partner has left call k. Anyone + // making that barrier one-sided silently breaks this. + int32_t* scratch_state = params.scratch_state; + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + const int call_epoch = advance_scratch_epoch(scratch_state, epoch_slot); + const int stage_offset = call_epoch * params.epoch_stride_packs; + // This rank owns the chunk at its own position in the island. + const int my_start = local * part; + const int my_end = (local == 3) ? params.num_packs : my_start + part; + const int my_slot = stage_offset + params.rank * params.rank_stride_packs; + const int cross_slot = stage_offset + cross_owner * params.rank_stride_packs; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + + // Island reduce-scatter. The contribution to this rank's own chunk stays + // on this GPU, so it costs no fabric time and needs no pass of its own. + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_u4_volatile(local_buffer, my_slot + idx, input[idx]); + } + for (int p = 0; p < 3; ++p) { + const int t = (local + 1 + p) & 3; + const int owner_t = base + t; + const int t_start = t * part; + const int t_end = (t == 3) ? params.num_packs : t_start + part; + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner_t]); + for (int idx = t_start + tid; idx < t_end; idx += stride) { + store_u4_volatile(owner_buffer, my_slot + idx, input[idx]); + } + __threadfence_system(); + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, p, flag, + island_mask); + } + + // Island sum, then the one cross-SYS exchange with the paired owner. + for (int idx = my_start + tid; idx < my_end; idx += stride) { + uint4 v0 = load_u4_volatile(local_buffer, + stage_offset + (base + 0) * params.rank_stride_packs + idx); + uint4 v1 = load_u4_volatile(local_buffer, + stage_offset + (base + 1) * params.rank_stride_packs + idx); + uint4 v2 = load_u4_volatile(local_buffer, + stage_offset + (base + 2) * params.rank_stride_packs + idx); + uint4 v3 = load_u4_volatile(local_buffer, + stage_offset + (base + 3) * params.rank_stride_packs + idx); + uint4 local_sum = packed_add_u4(packed_add_u4(v0, v1), packed_add_u4(v2, v3)); + // Keep the island sum so the gather phase does not recompute it. + store_u4_volatile(local_buffer, my_slot + idx, local_sum); + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + store_u4_volatile(cross_buffer, my_slot + idx, local_sum); + } + __threadfence_system(); + owner_pair_barrier(params.signal_ptrs, params.rank, params.rank, cross_owner, params.max_blocks, + 3, flag); + debug_cross_read_stall(params.rank); + + // Final value for the owned chunk, written locally. + for (int idx = my_start + tid; idx < my_end; idx += stride) { + uint4 mine = load_u4_volatile(local_buffer, my_slot + idx); + uint4 theirs = load_u4_volatile(local_buffer, cross_slot + idx); + uint4 final_value = packed_add_u4(mine, theirs); + reinterpret_cast(params.output)[idx] = final_value; + store_u4_volatile(local_buffer, my_slot + idx, final_value); + } + __threadfence_system(); + + // Island all-gather, staged on the same permutation schedule. + for (int p = 0; p < 3; ++p) { + const int peer = base + ((local + 1 + p) & 3); + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_u4_volatile(peer_buffer, my_slot + idx, + load_u4_volatile(local_buffer, my_slot + idx)); + } + __threadfence_system(); + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, 4 + p, flag, + island_mask); + } + + // Collect the three chunks owned by the other island members. + for (int p = 0; p < 3; ++p) { + const int t = (local + 1 + p) & 3; + const int owner_t = base + t; + const int t_start = t * part; + const int t_end = (t == 3) ? params.num_packs : t_start + part; + const int owner_slot = stage_offset + owner_t * params.rank_stride_packs; + for (int idx = t_start + tid; idx < t_end; idx += stride) { + reinterpret_cast(params.output)[idx] = + load_u4_volatile(local_buffer, owner_slot + idx); + } + } + // Hold the island until everyone has finished reading before the next + // call's reduce-scatter starts writing the same slots. This covers the + // intra-island reuse only; the cross-island edge is what the epoch double + // buffer above supplies. + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, 7, flag, island_mask); + } else { + Pack const* input = reinterpret_cast(params.input); + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_pack_volatile(local_buffer, my_slot + idx, input[idx]); + } + for (int p = 0; p < 3; ++p) { + const int t = (local + 1 + p) & 3; + const int owner_t = base + t; + const int t_start = t * part; + const int t_end = (t == 3) ? params.num_packs : t_start + part; + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner_t]); + for (int idx = t_start + tid; idx < t_end; idx += stride) { + store_pack_volatile(owner_buffer, my_slot + idx, input[idx]); + } + __threadfence_system(); + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, p, flag, + island_mask); + } + + for (int idx = my_start + tid; idx < my_end; idx += stride) { + Pack values[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + values[i] = load_pack_volatile( + local_buffer, stage_offset + (base + i) * params.rank_stride_packs + idx); + } + Pack local_sum = reduce_loaded_packs(values); + store_pack_volatile(local_buffer, my_slot + idx, local_sum); + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + store_pack_volatile(cross_buffer, my_slot + idx, local_sum); + } + __threadfence_system(); + owner_pair_barrier(params.signal_ptrs, params.rank, params.rank, cross_owner, params.max_blocks, + 3, flag); + debug_cross_read_stall(params.rank); + + for (int idx = my_start + tid; idx < my_end; idx += stride) { + Pack pair[2]; + pair[0] = load_pack_volatile(local_buffer, my_slot + idx); + pair[1] = load_pack_volatile(local_buffer, cross_slot + idx); + Pack final_value = reduce_loaded_packs(pair); + reinterpret_cast(params.output)[idx] = final_value; + store_pack_volatile(local_buffer, my_slot + idx, final_value); + } + __threadfence_system(); + + for (int p = 0; p < 3; ++p) { + const int peer = base + ((local + 1 + p) & 3); + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + for (int idx = my_start + tid; idx < my_end; idx += stride) { + store_pack_volatile(peer_buffer, my_slot + idx, + load_pack_volatile(local_buffer, my_slot + idx)); + } + __threadfence_system(); + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, 4 + p, flag, + island_mask); + } + + for (int p = 0; p < 3; ++p) { + const int t = (local + 1 + p) & 3; + const int owner_t = base + t; + const int t_start = t * part; + const int t_end = (t == 3) ? params.num_packs : t_start + part; + const int owner_slot = stage_offset + owner_t * params.rank_stride_packs; + for (int idx = t_start + tid; idx < t_end; idx += stride) { + reinterpret_cast(params.output)[idx] = + load_pack_volatile(local_buffer, owner_slot + idx); + } + } + block_barrier_mask(params.signal_ptrs, params.rank, 8, params.max_blocks, 7, flag, island_mask); + } + + if (threadIdx.x == 0) { + store_release_i32(self_signal + flag_offset(blockIdx.x, params.max_blocks, 8), flag); + } + // The barrier flag above must stay ahead of the release: a dependent kernel + // started by the trigger would otherwise observe this call's generation as + // not yet published. (The epoch is no longer a concern here -- it is + // committed at entry.) + debug_commit_per_block_epoch(epoch_slot, call_epoch); + pdl_grid_release_const(); +} + +template +__global__ __launch_bounds__(1024, 1) void ipc_topo_rsag8_block_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + pdl_grid_sync_const(); + + int32_t* self_signal = reinterpret_cast(params.signal_ptrs[params.rank]); + // Unsigned arithmetic for the bump: signed overflow is UB, and this counter + // is meant to wrap. generation_reached() reads it back on the circle. + int flag = + static_cast(static_cast(load_acquire_i32( + self_signal + flag_offset(blockIdx.x, params.max_blocks, 8))) + + 1u); + + // The caller's `blocks % 4 == 0` check is what keeps `blocks_per_chunk` + // non-zero; a zero stride below never advances the grid-stride loops. + int chunk = blockIdx.x & 3; + int chunk_block = blockIdx.x >> 2; + int blocks_per_chunk = gridDim.x >> 2; + int tid = chunk_block * blockDim.x + threadIdx.x; + int stride = blocks_per_chunk * blockDim.x; + int part = params.num_packs >> 2; + int start = chunk * part; + int end = (chunk == 3) ? params.num_packs : start + part; + int base = params.rank < 4 ? 0 : 4; + int owner = base + chunk; + int cross_owner = owner ^ 4; + // Call-level double buffer, for the same reason as the ring kernel: the + // phase 4 ack only covers this island, so nothing orders this rank's cross + // read against the paired owner's next-call cross write. Both TP8 kernels + // share this counter on purpose -- which one runs changes with the payload, + // and a counter advanced by only one of them would let a call land on + // a half two calls old. + int32_t* scratch_state = params.scratch_state; + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + const int call_epoch = advance_scratch_epoch(scratch_state, epoch_slot); + const int stage_offset = call_epoch * params.epoch_stride_packs; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + for (int idx = start + tid; idx < end; idx += stride) { + int offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(owner_buffer, offset, input[idx]); + } + __threadfence_system(); + island_owner_gather(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 1, flag); + + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + if (params.rank == owner) { + for (int idx = start + tid; idx < end; idx += stride) { + uint4 v0 = load_u4_volatile(local_buffer, + stage_offset + (base + 0) * params.rank_stride_packs + idx); + uint4 v1 = load_u4_volatile(local_buffer, + stage_offset + (base + 1) * params.rank_stride_packs + idx); + uint4 v2 = load_u4_volatile(local_buffer, + stage_offset + (base + 2) * params.rank_stride_packs + idx); + uint4 v3 = load_u4_volatile(local_buffer, + stage_offset + (base + 3) * params.rank_stride_packs + idx); + uint4 local_sum = packed_add_u4(packed_add_u4(v0, v1), packed_add_u4(v2, v3)); + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + int cross_write = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(cross_buffer, cross_write, local_sum); + } + } + if (params.rank == owner) { + __threadfence_system(); + } + owner_pair_barrier(params.signal_ptrs, params.rank, owner, cross_owner, params.max_blocks, 2, + flag); + debug_cross_read_stall(params.rank); + + if (params.rank == owner) { + for (int idx = start + tid; idx < end; idx += stride) { + uint4 v0 = load_u4_volatile(local_buffer, + stage_offset + (base + 0) * params.rank_stride_packs + idx); + uint4 v1 = load_u4_volatile(local_buffer, + stage_offset + (base + 1) * params.rank_stride_packs + idx); + uint4 v2 = load_u4_volatile(local_buffer, + stage_offset + (base + 2) * params.rank_stride_packs + idx); + uint4 v3 = load_u4_volatile(local_buffer, + stage_offset + (base + 3) * params.rank_stride_packs + idx); + uint4 local_sum = packed_add_u4(packed_add_u4(v0, v1), packed_add_u4(v2, v3)); + uint4 cross_sum = load_u4_volatile( + local_buffer, stage_offset + cross_owner * params.rank_stride_packs + idx); + uint4 final_value = packed_add_u4(local_sum, cross_sum); + reinterpret_cast(params.output)[idx] = final_value; +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + if (peer == params.rank) { + continue; + } + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(peer_buffer, final_offset, final_value); + } + } + } + if (params.rank == owner) { + __threadfence_system(); + } + island_owner_ready(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 3, flag); + + if (params.rank != owner) { + for (int idx = start + tid; idx < end; idx += stride) { + uint4 final_value = + load_u4_volatile(local_buffer, stage_offset + owner * params.rank_stride_packs + idx); + reinterpret_cast(params.output)[idx] = final_value; + } + } + // WRONG ORDER, and the reason enable_pdl is still refused at the binding: + // the trigger fires before island_owner_ack and before the barrier flag + // below, so a dependent kernel can start while this call's phase-4 ack and + // flag are still being written. Moving the release past both is the fix, + // but re-enabling PDL needs a per-kernel audit and an SM90 regression, so + // it is left visible rather than quietly reordered. + pdl_grid_release_const(); + island_owner_ack(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 4, flag); + } else { + Pack const* input = reinterpret_cast(params.input); + auto* owner_buffer = reinterpret_cast(params.tmp_ptrs[owner]); + for (int idx = start + tid; idx < end; idx += stride) { + int offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(owner_buffer, offset, input[idx]); + } + __threadfence_system(); + island_owner_gather(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 1, flag); + + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + if (params.rank == owner) { + for (int idx = start + tid; idx < end; idx += stride) { + Pack v0 = load_pack_volatile(local_buffer, + stage_offset + (base + 0) * params.rank_stride_packs + idx); + Pack v1 = load_pack_volatile(local_buffer, + stage_offset + (base + 1) * params.rank_stride_packs + idx); + Pack v2 = load_pack_volatile(local_buffer, + stage_offset + (base + 2) * params.rank_stride_packs + idx); + Pack v3 = load_pack_volatile(local_buffer, + stage_offset + (base + 3) * params.rank_stride_packs + idx); + Pack local_sum = add_pack(add_pack(v0, v1), add_pack(v2, v3)); + auto* cross_buffer = reinterpret_cast(params.tmp_ptrs[cross_owner]); + int cross_write = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(cross_buffer, cross_write, local_sum); + } + } + if (params.rank == owner) { + __threadfence_system(); + } + owner_pair_barrier(params.signal_ptrs, params.rank, owner, cross_owner, params.max_blocks, 2, + flag); + debug_cross_read_stall(params.rank); + + if (params.rank == owner) { + for (int idx = start + tid; idx < end; idx += stride) { + Pack v0 = load_pack_volatile(local_buffer, + stage_offset + (base + 0) * params.rank_stride_packs + idx); + Pack v1 = load_pack_volatile(local_buffer, + stage_offset + (base + 1) * params.rank_stride_packs + idx); + Pack v2 = load_pack_volatile(local_buffer, + stage_offset + (base + 2) * params.rank_stride_packs + idx); + Pack v3 = load_pack_volatile(local_buffer, + stage_offset + (base + 3) * params.rank_stride_packs + idx); + Pack local_sum = add_pack(add_pack(v0, v1), add_pack(v2, v3)); + Pack cross_sum = load_pack_volatile( + local_buffer, stage_offset + cross_owner * params.rank_stride_packs + idx); + Pack final_value = add_pack(local_sum, cross_sum); + reinterpret_cast(params.output)[idx] = final_value; +#pragma unroll + for (int peer_local = 0; peer_local < 4; ++peer_local) { + int peer = base + peer_local; + if (peer == params.rank) { + continue; + } + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int final_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(peer_buffer, final_offset, final_value); + } + } + } + if (params.rank == owner) { + __threadfence_system(); + } + island_owner_ready(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 3, flag); + + if (params.rank != owner) { + for (int idx = start + tid; idx < end; idx += stride) { + Pack final_value = load_pack_volatile( + local_buffer, stage_offset + owner * params.rank_stride_packs + idx); + reinterpret_cast(params.output)[idx] = final_value; + } + } + // WRONG ORDER, and the reason enable_pdl is still refused at the binding: + // the trigger fires before island_owner_ack and before the barrier flag + // below, so a dependent kernel can start while this call's phase-4 ack and + // flag are still being written. Moving the release past both is the fix, + // but re-enabling PDL needs a per-kernel audit and an SM90 regression, so + // it is left visible rather than quietly reordered. + pdl_grid_release_const(); + island_owner_ack(params.signal_ptrs, params.rank, base, owner, params.max_blocks, 4, flag); + } + + if (threadIdx.x == 0) { + store_release_i32(self_signal + flag_offset(blockIdx.x, params.max_blocks, 8), flag); + } + debug_commit_per_block_epoch(epoch_slot, call_epoch); +} + +template +__global__ __launch_bounds__(1024, 1) void push_oneshot_param_kernel( + const PushOneshotParamData __grid_constant__ params) { + using Pack = typename PackTraits::Pack; + pdl_grid_sync_const(); + + int32_t* epoch_slot = + params.epoch_slots + blockIdx.x; // used only by the PER_BLOCK_EPOCH debug build + int epoch = advance_scratch_epoch(params.scratch_state, epoch_slot); + int stage_offset = epoch * params.epoch_stride_packs; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + + if constexpr (std::is_same_v || std::is_same_v) { + uint4 const* input = reinterpret_cast(params.input); + for (int idx = tid; idx < params.num_packs; idx += stride) { + uint4 value = clear_pos_zero_u4_16(input[idx]); +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + auto* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int peer_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_u4_volatile(peer_buffer, peer_offset, value); + } + } + + auto* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + uint4 reset = {0u, 0u, 0u, 0u}; + for (int idx = tid; idx < params.num_packs; idx += stride) { + uint4 values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int peer_offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_u4_volatile(local_buffer, peer_offset); + waiting |= has_pos_zero_u4_16(values[peer]); + } + if (!waiting) { + break; + } + } + + uint4 acc; + if constexpr (Fp32Reduce) { + acc = reduce_u4_fp32(values); + } else { + acc = values[0]; +#pragma unroll + for (int peer = 1; peer < WorldSize; ++peer) { + acc = packed_add_u4(acc, values[peer]); + } + } + reinterpret_cast(params.output)[idx] = acc; + +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int peer_offset = stage_offset + peer * params.rank_stride_packs + idx; + local_buffer[peer_offset] = reset; + } + } + } else { + Pack const* input = reinterpret_cast(params.input); + for (int idx = tid; idx < params.num_packs; idx += stride) { + Pack value = input[idx]; + clear_pos_zero_pack(value); +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + Pack* peer_buffer = reinterpret_cast(params.tmp_ptrs[peer]); + int peer_offset = stage_offset + params.rank * params.rank_stride_packs + idx; + store_pack_volatile(peer_buffer, peer_offset, value); + } + } + + Pack* local_buffer = reinterpret_cast(params.tmp_ptrs[params.rank]); + Pack reset = zero_pack(); + for (int idx = tid; idx < params.num_packs; idx += stride) { + Pack values[WorldSize]; + while (true) { + bool waiting = false; +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int peer_offset = stage_offset + peer * params.rank_stride_packs + idx; + values[peer] = load_pack_volatile(local_buffer, peer_offset); + waiting |= has_pos_zero_pack(values[peer]); + } + if (!waiting) { + break; + } + } + + Pack acc = reduce_loaded_packs(values); + reinterpret_cast(params.output)[idx] = acc; + +#pragma unroll + for (int peer = 0; peer < WorldSize; ++peer) { + int peer_offset = stage_offset + peer * params.rank_stride_packs + idx; + local_buffer[peer_offset] = reset; + } + } + } + + debug_commit_per_block_epoch(epoch_slot, epoch); + pdl_grid_release_const(); +} +// --------------------------------------------------------------------------- +// Host side +// --------------------------------------------------------------------------- + +// Byte layout of one rank's workspace slab. +// +// [ epoch slots | barrier phase slots | barrier flags +// | block-scratch epoch + arrival | pack scratch | block scratch ] +// +// Both scratch regions are sized for world_size ranks x a double-buffered +// epoch, so a rank may start collective N+1 before its peer has drained N. +// The epoch halves sit at fixed offsets derived from max_numel rather than +// from the current payload: if they moved with the payload, a rank that +// finished a large collective and flipped its epoch would start writing a +// small one inside the region a lagging peer is still draining. +struct WorkspaceLayout { + size_t signal_bytes; + size_t max_payload_bytes; + size_t scratch_bytes; // per scratch region + size_t total_bytes; +}; + +inline WorkspaceLayout compute_workspace_layout(int world_size, int64_t max_numel, int elem_size, + int max_blocks) { + const size_t epoch_slots = static_cast(max_blocks); + const size_t barrier_slots = static_cast(kSignalPhases) * + static_cast(max_blocks) * static_cast(world_size); + const size_t flag_slots = static_cast(max_blocks); + // {epoch, arrival} per scratch region, in ScratchRegion order. Appended at + // the tail so phase_offset() and flag_offset(), both anchored at the front, + // are unchanged. See scratch_state_offset(). + const size_t scratch_state_slots = 2 * 2; + const size_t signal_slots = epoch_slots + barrier_slots + flag_slots + scratch_state_slots; + auto align128 = [](size_t n) { return (n + 127u) & ~static_cast(127u); }; + WorkspaceLayout layout{}; + layout.signal_bytes = align128(sizeof(int32_t) * signal_slots); + layout.max_payload_bytes = + align128(static_cast(max_numel) * static_cast(elem_size)); + layout.scratch_bytes = align128(2 * static_cast(world_size) * layout.max_payload_bytes); + layout.total_bytes = layout.signal_bytes + 2 * layout.scratch_bytes; + return layout; +} + +// Bytes each rank must allocate and share over CUDA IPC. +inline int64_t workspace_size(int world_size, int64_t max_numel, int elem_size, int max_blocks) { + return static_cast( + compute_workspace_layout(world_size, max_numel, elem_size, max_blocks).total_bytes); +} + +// Per-region device pointers into every rank's slab, as seen by this process. +struct PeerViews { + uint64_t signal[kMaxWorldSize]; + uint64_t pack[kMaxWorldSize]; + uint64_t block[kMaxWorldSize]; + int32_t* self_signal; +}; + +// ipc_ptrs[i] must address rank i's slab; ipc_ptrs[rank] is this rank's own. +inline PeerViews make_peer_views(const int64_t* ipc_ptrs, int world_size, int rank, + const WorkspaceLayout& layout) { + PeerViews views{}; + for (int peer = 0; peer < world_size; ++peer) { + auto* base = reinterpret_cast(ipc_ptrs[peer]); + views.signal[peer] = reinterpret_cast(base); + auto* scratch = base + static_cast(layout.signal_bytes); + views.pack[peer] = reinterpret_cast(scratch); + views.block[peer] = reinterpret_cast(scratch + layout.scratch_bytes); + } + views.self_signal = reinterpret_cast(ipc_ptrs[rank]); + return views; +} + +template +inline cudaError_t launch(Kernel kernel, dim3 grid, dim3 block, cudaStream_t stream, bool use_pdl, + Args const&... args) { +#if CUDART_VERSION >= 12000 + if (use_pdl) { + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0; + config.stream = stream; + config.attrs = attr; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, args...); + } +#else + if (use_pdl) return cudaErrorNotSupported; +#endif + kernel<<>>(args...); + return cudaGetLastError(); +} + +// Kernel selection. `variant` is picked by the caller rather than by a +// threshold here, because the crossovers depend on the fabric and are measured +// per machine. +// +// world variant kernel +// 2 kUnstaged ipc_tp2_remote_push_kernel (block scratch) +// 2 kStaged ipc_tp2_remote_push_kernel (block scratch) +// 4 kUnstaged push_oneshot_param_kernel (pack scratch) +// 4 kStaged ipc_rsag_push_param_kernel<4> (block scratch) +// 4 kStagedRing ipc_rsag_ring_push_param_kernel<4> (block scratch) +// 8 kUnstaged ipc_topo_rsag8_push_param_kernel (pack scratch) +// 8 kStaged ipc_topo_rsag8_block_param_kernel (block scratch) +// 8 kStagedRing ipc_topo_rsag8_ring_push_param_kernel (block scratch) +// 8 kFlatStaged ipc_rsag_push_param_kernel<8> (pack scratch) +// +// Preconditions the caller must have validated: world_size in {2,4,8}; the +// (world_size, variant) pair appears above; 0 < blocks <= max_blocks; +// 0 < threads <= 1024; numel and max_numel both divisible by the 16-byte pack +// width; numel * elem_size <= max_payload_bytes; and blocks % 4 == 0 for +// (8, kStaged), since that kernel derives its chunk from blockIdx.x & 3. +template +cudaError_t all_reduce(const T* input, T* output, int64_t numel, const PeerViews& views, int rank, + int world_size, int max_blocks, int64_t max_numel, int blocks, int threads, + Variant variant, bool use_pdl, cudaStream_t stream) { + using Traits = PackTraits; + const int num_packs = static_cast(numel / Traits::kPackElems); + const int rank_stride_packs = static_cast(max_numel / Traits::kPackElems); + const dim3 grid(static_cast(blocks)); + const dim3 cta(static_cast(threads)); + + if (world_size == 2) { + IpcTp2RemotePushData params{}; + params.tmp_ptrs[0] = views.block[0]; + params.tmp_ptrs[1] = views.block[1]; + params.input = input; + params.output = output; + params.epoch_slots = views.self_signal; + // TP2 stages through views.block under either variant, so it shares the + // block region's counter -- nominal here, since it is the only TP2 kernel, + // but the state must follow the region it actually writes. + params.scratch_state = + views.self_signal + scratch_state_offset(max_blocks, 2, ScratchRegion::kBlock); + params.num_packs = num_packs; + params.rank_stride_packs = rank_stride_packs; + params.rank = rank; + const bool staged = variant == Variant::kStaged; + if (use_pdl) { + return staged ? launch(ipc_tp2_remote_push_kernel, grid, cta, stream, true, + params) + : launch(ipc_tp2_remote_push_kernel, grid, cta, stream, true, + params); + } + return staged ? launch(ipc_tp2_remote_push_kernel, grid, cta, stream, false, + params) + : launch(ipc_tp2_remote_push_kernel, grid, cta, stream, false, + params); + } + + PushOneshotParamData params{}; + // Region and counter are chosen together: a kernel reading one region while + // advancing another's epoch would corrupt both. Which kernels may share a + // region is a protocol question, not a partitioning one -- see ScratchRegion. + const ScratchRegion region = (variant == Variant::kUnstaged || variant == Variant::kFlatStaged) + ? ScratchRegion::kPack + : ScratchRegion::kBlock; + const uint64_t* scratch = region == ScratchRegion::kPack ? views.pack : views.block; + for (int peer = 0; peer < world_size; ++peer) { + params.tmp_ptrs[peer] = scratch[peer]; + params.signal_ptrs[peer] = views.signal[peer]; + } + params.input = input; + params.output = output; + params.epoch_slots = views.self_signal; + params.scratch_state = views.self_signal + scratch_state_offset(max_blocks, world_size, region); + params.num_packs = num_packs; + params.rank_stride_packs = rank_stride_packs; + params.epoch_stride_packs = world_size * rank_stride_packs; + params.rank = rank; + params.max_blocks = max_blocks; + +#define FI_PCIE_IPC_LAUNCH(KERNEL_EXPR, PDL) launch(KERNEL_EXPR, grid, cta, stream, PDL, params) + +#define FI_PCIE_IPC_SELECT(PDL) \ + do { \ + if (world_size == 8) { \ + switch (variant) { \ + case Variant::kUnstaged: \ + return FI_PCIE_IPC_LAUNCH((ipc_topo_rsag8_push_param_kernel), PDL); \ + case Variant::kStaged: \ + return FI_PCIE_IPC_LAUNCH((ipc_topo_rsag8_block_param_kernel), PDL); \ + case Variant::kStagedRing: \ + return FI_PCIE_IPC_LAUNCH((ipc_topo_rsag8_ring_push_param_kernel), PDL); \ + case Variant::kFlatStaged: \ + return FI_PCIE_IPC_LAUNCH((ipc_rsag_push_param_kernel), PDL); \ + } \ + return cudaErrorInvalidValue; \ + } \ + switch (variant) { \ + case Variant::kUnstaged: \ + return FI_PCIE_IPC_LAUNCH((push_oneshot_param_kernel), PDL); \ + case Variant::kStaged: \ + return FI_PCIE_IPC_LAUNCH((ipc_rsag_push_param_kernel), PDL); \ + case Variant::kStagedRing: \ + return FI_PCIE_IPC_LAUNCH((ipc_rsag_ring_push_param_kernel), PDL); \ + default: \ + return cudaErrorInvalidValue; \ + } \ + } while (false) + + if (use_pdl) { + FI_PCIE_IPC_SELECT(true); + } + FI_PCIE_IPC_SELECT(false); + +#undef FI_PCIE_IPC_SELECT +#undef FI_PCIE_IPC_LAUNCH +} + +} // namespace pcie_ipc +} // namespace comm +} // namespace flashinfer + +#endif // FLASHINFER_COMM_PCIE_IPC_ALL_REDUCE_CUH_ diff --git a/scripts/task_test_single_node_comm_kernels.sh b/scripts/task_test_single_node_comm_kernels.sh index 071a4e45cb4..daea0d8a1e6 100644 --- a/scripts/task_test_single_node_comm_kernels.sh +++ b/scripts/task_test_single_node_comm_kernels.sh @@ -28,6 +28,20 @@ python -c "import nvshmem.core" 2>/dev/null || pip install --no-deps nvshmem4py- # vllm ar pytest -s tests/comm/test_vllm_custom_allreduce.py +# pcie ipc ar (intra-node PCIe without NVLink) +# Every case here needs 2, 4 or 8 GPUs. With this script's default +# CUDA_VISIBLE_DEVICES=0 they all skip and pytest still exits 0, which reads as +# "passed" in a log. Say so loudly rather than letting a green run mean nothing. +pcie_ipc_gpus=$(python3 -c 'import torch; print(torch.cuda.device_count())') +if [ "$pcie_ipc_gpus" -lt 2 ]; then + echo "############################################################" + echo "# SKIPPING pcie ipc ar: $pcie_ipc_gpus GPU visible, needs >=2 (>=8 for full" + echo "# coverage). This is NOT a pass. Set CUDA_VISIBLE_DEVICES=0,1,...,7." + echo "############################################################" +else + echo "pcie ipc ar: $pcie_ipc_gpus GPUs visible (8 needed for full coverage)" + pytest -s tests/comm/test_pcie_ipc_all_reduce.py +fi # trtllm ar + fusion pytest -s tests/comm/test_trtllm_allreduce.py pytest -s tests/comm/test_trtllm_allreduce_fusion.py diff --git a/tests/comm/test_pcie_ipc_all_reduce.py b/tests/comm/test_pcie_ipc_all_reduce.py new file mode 100644 index 00000000000..ea895600fcb --- /dev/null +++ b/tests/comm/test_pcie_ipc_all_reduce.py @@ -0,0 +1,1373 @@ +""" +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. +""" + +import multiprocessing as mp +import os +import socket +import time +import warnings +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +import flashinfer.comm as comm +from flashinfer.comm.pcie_ipc_policy import IpcLaunchConfig, IpcVariant +from flashinfer.comm.pcie_ipc_topology import PCIE_IPC_PROFILES + +# (world_size, hidden, batch, blocks, threads, variant) +# +# Between them these cases select every kernel the header can dispatch to. +# Spelled out rather than resolved: the seed reaches five of the nine (world +# size, variant) pairs, and the other four -- (2, STAGED), (4, UNSTAGED), +# (8, STAGED), (8, FLAT_STAGED) -- are reachable only through an explicit +# config or the tuner. +_JOIN_TIMEOUT_S = 600 + +_CASES = [ + (2, 2048, 1, 32, 64, IpcVariant.UNSTAGED), # tp2, unstaged + (2, 2048, 8, 96, 64, IpcVariant.STAGED), # tp2, staged + (2, 2048, 128, 16, 128, IpcVariant.UNSTAGED), + (4, 4096, 1, 1, 128, IpcVariant.STAGED), # rsag push + (4, 4096, 8, 1, 256, IpcVariant.STAGED_RING), # rsag ring push + (4, 4096, 128, 4, 256, IpcVariant.STAGED_RING), + (4, 4096, 16, 8, 256, IpcVariant.UNSTAGED), # one-shot push at 4 ranks + (8, 6144, 1, 32, 128, IpcVariant.UNSTAGED), # topo pack + (8, 6144, 2, 1, 128, IpcVariant.FLAT_STAGED), # generic rsag at 8 ranks + (8, 6144, 8, 1, 256, IpcVariant.STAGED_RING), # topo ring push + (8, 6144, 128, 2, 256, IpcVariant.STAGED_RING), + (8, 6144, 16, 8, 256, IpcVariant.STAGED), # topo block (blocks % 4 == 0) +] + + +def get_open_port() -> int: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + except OSError: + with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _init_process_group(world_size: int, rank: int, port: int) -> None: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + ) + + +def multi_process_parallel( + world_size: int, + target: Any, + args: tuple = (), + timeout_s: float = _JOIN_TIMEOUT_S, +) -> None: + """Run ``target`` on ``world_size`` ranks, or fail with a bounded wait. + + ``timeout_s`` is lowered by the negative-control tests, which *expect* the + deadlock and would otherwise pay the full timeout to observe it. + """ + mp.set_start_method("spawn", force=True) + port = get_open_port() + procs = [] + for rank in range(world_size): + p = mp.Process( + target=target, args=(world_size, rank, port) + args, name=f"Worker-{rank}" + ) + p.start() + procs.append(p) + # A bounded join is essential here: the failure mode this suite is built to + # catch is a spin-wait deadlock, and an unbounded join would hang CI rather + # than report it. + # + # Every exit runs the reaper, not just the timeout one. When one rank dies + # the others are left spinning inside a collective their peer will never + # reach, so returning without them would leak processes that hold a GPU + # context for the rest of the session -- and the next test in the file + # would then fail for reasons that have nothing to do with it. + deadline = time.monotonic() + timeout_s + try: + for rank, p in enumerate(procs): + p.join(timeout=max(1.0, deadline - time.monotonic())) + if p.is_alive(): + raise AssertionError( + f"rank {rank} did not finish within {timeout_s}s; " + "the collective most likely deadlocked" + ) + assert p.exitcode == 0, f"rank {rank} failed with exit code {p.exitcode}" + finally: + for p in procs: + if p.is_alive(): + p.kill() + for p in procs: + p.join(timeout=30.0) + + +def _correctness_worker( + world_size: int, rank: int, port: int, dtype: torch.dtype +) -> None: + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + cases = [c for c in _CASES if c[0] == world_size] + max_numel = max(h * b for _, h, b, *_ in cases) + + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=max_numel, dtype=dtype, max_blocks=128 + ) + # The profile only keys the tune cache, but the probe must still land + # on a known one; which one is a property of the machine. + assert ws.profile in PCIE_IPC_PROFILES, ws.profile_reason + + for _, hidden, batch, blocks, threads, variant in cases: + shape = (batch, hidden) + # Small integers so the reduction is exact in every dtype and the + # comparison can use a zero tolerance; the kernel sums in a + # different order than NCCL, which would otherwise show up. + inp = torch.randint(0, 16, shape, dtype=torch.int32, device=device).to( + dtype + ) + # Snapshot before any call, and build the reference from the + # snapshot: cloning after the fact would follow the input if the + # kernel mutated it, hiding exactly what we mean to check. + inp_before = inp.clone() + ref = inp_before.clone() + dist.all_reduce(ref, group=group) + + # Admitted shapes must resolve and be correct, not merely the + # right shape: a mis-selected kernel would otherwise go unnoticed. + assert ws.supports(inp) + torch.testing.assert_close(ws.all_reduce(inp), ref, rtol=0, atol=0) + + # The explicit config then pins the kernel this case exists for. + out = ws.all_reduce( + inp, + config=IpcLaunchConfig(blocks, threads, variant), + ) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + # The sentinel rewrite must happen on a register copy, never in the + # caller's buffer. + assert torch.equal(inp, inp_before) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _unsupported_shape_worker(world_size: int, rank: int, port: int) -> None: + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=8192, dtype=torch.bfloat16 + ) + + # Hidden size is not an admission criterion. Only the byte count is, and + # this one fits. + assert ws.supports(torch.empty(2, 4096, dtype=torch.bfloat16, device=device)) + # Fewer 16-byte packs than ranks: the reduce-scatter split would hand + # the whole payload to one owner and leave the others idle. + assert not ws.supports(torch.empty(8, dtype=torch.bfloat16, device=device)) + # Larger than the workspace. + assert not ws.supports(torch.empty(16384, dtype=torch.bfloat16, device=device)) + # Element size does not match the workspace. + assert not ws.supports(torch.empty(1024, dtype=torch.float32, device=device)) + # Byte size not a multiple of 16. + assert not ws.supports(torch.empty(4, dtype=torch.bfloat16, device=device)) + # Non-contiguous. + assert not ws.supports( + torch.empty(64, 64, dtype=torch.bfloat16, device=device).t() + ) + + # A device mismatch is a caller bug, so it raises instead of reporting + # "unsupported". Reporting unsupported would send this rank to another + # backend while its peers stay here, and the group would hang. + peer = torch.device(f"cuda:{(rank + 1) % world_size}") + for wrong in (torch.device("cpu"), peer): + with pytest.raises(ValueError, match="workspace was built on"): + ws.supports(torch.empty(1024, dtype=torch.bfloat16, device=wrong)) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _cuda_graph_worker(world_size: int, rank: int, port: int) -> None: + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = {2: 2048, 4: 4096, 8: 6144}[world_size] + batch = 8 + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * batch, dtype=torch.bfloat16 + ) + inp = torch.randint( + 0, 16, (batch, hidden), dtype=torch.int32, device=device + ).to(torch.bfloat16) + out = torch.empty_like(inp) + + # Warm up outside the graph. The protocol state lives in device memory + # and is read at kernel entry, so replays pick up whatever the previous + # launch left; nothing is baked into the captured node. + for _ in range(3): + ws.all_reduce(inp, out=out) + torch.cuda.synchronize() + dist.barrier(group=group) + + ref = inp.clone() + dist.all_reduce(ref, group=group) + + # Capture an odd count as well as an even one. Protocol state that + # alternates per call returns to its starting parity after an even + # capture, so only an odd capture exercises a replay that begins on the + # opposite parity from the one before it. + for captured in (3, 4): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(captured): + ws.all_reduce(inp, out=out) + for _ in range(16): + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, ref, rtol=0, atol=0) + dist.barrier(group=group) + del graph + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _shape_change_worker(world_size: int, rank: int, port: int) -> None: + """Interleave shapes with no intervening synchronisation. + + The epoch double buffer lets a rank start the next collective before its + peer has drained the previous one, so a shape change is exactly when a + staging offset derived from the current payload rather than from the + workspace capacity would corrupt. + + Consecutive calls must carry *different* payloads. A reuse-too-early bug + overwrites a scratch slot with another call's copy of the same reduction; + if every call reduces the same tensor the two values are bit-identical and + the corruption is invisible no matter how the calls are interleaved. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = {2: 2048, 4: 4096, 8: 6144}[world_size] + batches = [1, 16, 1, 96, 1] + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * max(batches), dtype=torch.bfloat16 + ) + # Four payloads per shape, cycled, so a call never repeats the data of + # the call two before it -- the distance at which the scratch half is + # reused. + variants = 4 + inputs = { + b: [ + torch.randint(0, 16, (b, hidden), dtype=torch.int32, device=device).to( + torch.bfloat16 + ) + for _ in range(variants) + ] + for b in batches + } + refs = {} + for b, xs in inputs.items(): + refs[b] = [] + for x in xs: + r = x.clone() + dist.all_reduce(r, group=group) + refs[b].append(r) + dist.barrier(group=group) + + outs = [] + for i in range(40): + for b in batches: + v = i % variants + outs.append((b, v, ws.all_reduce(inputs[b][v]))) + torch.cuda.synchronize() + for b, v, o in outs: + torch.testing.assert_close(o, refs[b][v], rtol=0, atol=0) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_pcie_ipc_cuda_graph(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _cuda_graph_worker) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_pcie_ipc_shape_change(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _shape_change_worker) + + +# Batches straddling a seed block-count boundary at 4 ranks, hidden 4096: 32 -> +# one ring block, 64/70/80 -> two. Each pair grows the grid; the last two share +# one and are the control. Variant and thread count are constant, so the grid is +# the only thing that moves. +_GRID_CHANGE_BATCHES = [32, 64, 32, 70, 32, 80, 70] + + +def _grid_change_worker(world_size: int, rank: int, port: int) -> None: + """Change the grid size between calls, which used to be an instant hang. + + The epoch used to be per block: each block flipped its own parity on exit, + so the parity recorded how many times *that block* had run. Grow the grid + and a block running for the first time reads the initial 0 and picks the + half the previous call just used, with no slack at all. + + The sequence ends with ``80, 70`` deliberately. Those two calls have the + same grid, so nothing about *them* is unusual -- but the earlier growth has + already desynchronised the block ranges, and under the old scheme they hung + anyway. One grid change was enough to arm it permanently, which is why + testing only at the boundary would not be enough. + + Deterministic: no stall, no instrumentation, no race to lose. The failure + was a hang rather than a wrong answer, because in this kernel the racing + write lands on the address the victim polls next. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = 4096 + # No profile pinned: the seed does not consult it, so the batch->blocks + # mapping above holds on any machine. + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=hidden * max(_GRID_CHANGE_BATCHES), + dtype=torch.bfloat16, + ) + # Without a grid change there is no first-appearance block and nothing + # under test, so the sequence is checked rather than assumed. + grids = { + ws.launch_config( + torch.empty(b, hidden, dtype=torch.bfloat16, device=device) + ).blocks + for b in _GRID_CHANGE_BATCHES + } + assert len(grids) > 1, f"every batch resolves to the same grid: {grids}" + # Distinct payloads: a repeated one makes a mis-addressed write store a + # bit-identical value, hiding any corruption that does not also hang. + inputs = [ + torch.randint(1, 16, (b, hidden), dtype=torch.int32, device=device).to( + torch.bfloat16 + ) + + i + for i, b in enumerate(_GRID_CHANGE_BATCHES) + ] + refs = [] + for x in inputs: + r = x.clone() + dist.all_reduce(r, group=group) + refs.append(r) + dist.barrier(group=group) + + # Nothing drains the queue between calls -- that is the point. + outs = [ws.all_reduce(x) for x in inputs] + for out, ref, batch in zip(outs, refs, _GRID_CHANGE_BATCHES, strict=True): + torch.testing.assert_close( + out, ref, rtol=0, atol=0, msg=f"batch {batch} in the grid-change run" + ) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [4]) +def test_pcie_ipc_grid_change(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _grid_change_worker) + + +def _single_block_epoch_worker(world_size: int, rank: int, port: int) -> None: + """A one-block launch must leave the epoch exactly where a larger one would. + + ``gridDim.x == 1`` skips the arrival atomic -- the sole CTA is trivially the + last arrival -- and flips the epoch directly. That is only sound if the two + paths leave identical state, which is invisible from outside except through + what the *next* call reads. Alternating the two therefore checks it: if the + fast path left the epoch unflipped, or flipped twice, a following call would + land on the half its predecessor just used. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = 4096 + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * 8, dtype=torch.bfloat16, profile="rootcplx" + ) + one = IpcLaunchConfig(1, 256, IpcVariant.STAGED_RING) + many = IpcLaunchConfig(4, 256, IpcVariant.STAGED_RING) + payloads = [ + torch.randint(1, 16, (8, hidden), dtype=torch.int32, device=device).to( + torch.bfloat16 + ) + + i + for i in range(4) + ] + refs = [] + for x in payloads: + r = x.clone() + dist.all_reduce(r, group=group) + refs.append(r) + dist.barrier(group=group) + + # Every ordering of the two paths, back to back, nothing draining. + outs = [] + for cfg_a, cfg_b in ((one, many), (many, one), (one, one), (many, many)): + for cfg, x in zip((cfg_a, cfg_b), payloads[:2], strict=True): + outs.append((ws.all_reduce(x, config=cfg), cfg)) + for i, (out, cfg) in enumerate(outs): + torch.testing.assert_close( + out, refs[i % 2], rtol=0, atol=0, msg=f"call {i} at blocks={cfg.blocks}" + ) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [4]) +def test_pcie_ipc_single_block_epoch_matches(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _single_block_epoch_worker) + + +# One call per scratch region at 4 ranks: the staged kernels stage through +# kBlock, the one-shot push through kPack (see ScratchRegion in +# pcie_ipc_all_reduce.cuh). Paired with the batch each is issued at, so the +# region and the owner partition move together. +_REGION_ALTERNATION = [ + (96, IpcLaunchConfig(2, 256, IpcVariant.STAGED)), + (1, IpcLaunchConfig(8, 256, IpcVariant.UNSTAGED)), + (96, IpcLaunchConfig(2, 256, IpcVariant.STAGED)), + (1, IpcLaunchConfig(8, 256, IpcVariant.UNSTAGED)), + (1, IpcLaunchConfig(8, 256, IpcVariant.UNSTAGED)), + (96, IpcLaunchConfig(2, 256, IpcVariant.STAGED)), +] + + +def _region_alternation_worker(world_size: int, rank: int, port: int) -> None: + """Alternate the two scratch regions on one workspace. + + Each region carries its own call-level epoch counter, so only an + interleaved sequence drives them against each other -- and only such a + sequence catches a state pointer bound to the wrong region. A run that + stays in one region never reads the other's leftovers and passes either way. + + The configurations are explicit because nothing has to choose them: the + hazard is in the kernels, not in which kernel a shape resolves to. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = 4096 + # An edit that left every call in one region would still pass while + # checking nothing. + regions = {cfg.variant == IpcVariant.UNSTAGED for _, cfg in _REGION_ALTERNATION} + assert regions == {True, False}, "the sequence stays in one scratch region" + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=hidden * max(b for b, _ in _REGION_ALTERNATION), + dtype=torch.bfloat16, + ) + payloads, refs = [], [] + for i, (b, _) in enumerate(_REGION_ALTERNATION): + x = ( + torch.randint(1, 16, (b, hidden), dtype=torch.int32, device=device).to( + torch.bfloat16 + ) + + i + ) + r = x.clone() + dist.all_reduce(r, group=group) + payloads.append(x) + refs.append(r) + dist.barrier(group=group) + + for x, ref, (b, cfg) in zip(payloads, refs, _REGION_ALTERNATION, strict=True): + torch.testing.assert_close( + ws.all_reduce(x, config=cfg), + ref, + rtol=0, + atol=0, + msg=f"batch {b} with {cfg}", + ) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [4]) +def test_pcie_ipc_scratch_region_alternation(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _region_alternation_worker) + + +# Every TP8 variant at several launch configurations. The tuner can pick any of +# them for a neighbouring batch, so this is a sequence a caller can produce. +# +# The order is not arbitrary. Each call flips its region's epoch, so a sequence +# that alternates strictly between the sentinel and topology kernels locks each +# of them onto one half and the sentinel kernel never meets the other's +# leftovers -- it passes whatever region they are in. Two topology calls in a +# row dirty both halves, and two sentinel calls in a row then have to read both. +_MIXED_TP8 = [ + IpcLaunchConfig(32, 128, IpcVariant.UNSTAGED), + IpcLaunchConfig(4, 128, IpcVariant.STAGED), + IpcLaunchConfig(1, 128, IpcVariant.STAGED_RING), + IpcLaunchConfig(1, 128, IpcVariant.FLAT_STAGED), + IpcLaunchConfig(2, 256, IpcVariant.FLAT_STAGED), + IpcLaunchConfig(12, 256, IpcVariant.UNSTAGED), + IpcLaunchConfig(8, 256, IpcVariant.STAGED), + IpcLaunchConfig(1, 512, IpcVariant.FLAT_STAGED), +] +# The owner partition is derived from the pack count, so the batch has to move +# too: a fixed shape only ever exercises one set of chunk boundaries. +_MIXED_TP8_BATCHES = [1, 2, 3, 5, 8] + + +def _mixed_variant_worker(world_size: int, rank: int, port: int) -> None: + """Every TP8 variant interleaved on one workspace, across several shapes. + + Sentinel kernels and barrier kernels cannot share a scratch region (see + ScratchRegion in pcie_ipc_all_reduce.cuh). A violation is wrong output on a + subset of ranks, not a hang, which is why this asserts group-wide. + + Distinct payloads per call are essential. With one payload reused, a stale + read returns the value the call would have computed anyway and the whole + sequence passes while proving nothing. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = 6144 + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=hidden * max(_MIXED_TP8_BATCHES), + dtype=torch.bfloat16, + profile="rootcplx", + ) + call = 0 + for batch in _MIXED_TP8_BATCHES: + for config in _MIXED_TP8: + inp = torch.randint( + 0, 16, (batch, hidden), dtype=torch.int32, device=device + ).to(torch.bfloat16) + float(call % 11) + ref = inp.clone() + dist.all_reduce(ref, group=group) + out = ws.all_reduce(inp, config=config) + # Group-wide: the hazard corrupts a subset of ranks, so a + # rank-local assertion can pass on rank 0 while the collective + # is wrong elsewhere. + wrong = torch.tensor([int((out != ref).sum().item())], device=device) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, ( + f"call {call} (batch {batch}, {config}) produced " + f"{int(wrong.item())} wrong elements: a variant is reading " + "scratch another variant left dirty" + ) + call += 1 + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_mixed_variants_share_no_scratch(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _mixed_variant_worker) + + +# Legal (blocks, threads, variant) per world size, one entry per dispatchable +# kernel. blocks is the smallest each kernel accepts: the TP8 block kernel +# derives its chunk from blockIdx.x & 3 and strides by gridDim.x >> 2. +_TINY_CONFIGS = { + 2: [ + IpcLaunchConfig(1, 64, IpcVariant.UNSTAGED), + IpcLaunchConfig(1, 64, IpcVariant.STAGED), + ], + 4: [ + IpcLaunchConfig(1, 64, IpcVariant.UNSTAGED), + IpcLaunchConfig(1, 64, IpcVariant.STAGED), + IpcLaunchConfig(1, 64, IpcVariant.STAGED_RING), + ], + 8: [ + IpcLaunchConfig(1, 64, IpcVariant.UNSTAGED), + IpcLaunchConfig(4, 64, IpcVariant.STAGED), + IpcLaunchConfig(1, 64, IpcVariant.STAGED_RING), + IpcLaunchConfig(1, 64, IpcVariant.FLAT_STAGED), + ], +} + + +def _tiny_payload_worker(world_size: int, rank: int, port: int) -> None: + """Payloads with fewer 16-byte packs than ranks, on every variant. + + The reduce-scatter split gives each rank ``num_packs // world_size`` packs, + which is zero once the payload is smaller than one pack per rank. The + staged kernels have to agree on who owns the payload in that case; a kernel + that writes to one owner and polls another spins forever rather than + returning a wrong answer, so the bounded join in multi_process_parallel is + what turns a regression into a report. + + Reached through explicit configurations: admission requires at least one + pack per rank, so no shape the policy accepts lands here. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + pack_elems = 8 # 16 bytes / bfloat16 + numels = [pack_elems * n for n in range(1, world_size + 1)] + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=max(numels), + dtype=torch.bfloat16, + profile="rootcplx", + ) + for numel in numels: + for config in _TINY_CONFIGS[world_size]: + inp = torch.randint( + 0, 16, (1, numel), dtype=torch.int32, device=device + ).to(torch.bfloat16) + ref = inp.clone() + dist.all_reduce(ref, group=group) + out = ws.all_reduce(inp, config=config) + wrong = torch.tensor([int((out != ref).sum().item())], device=device) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, ( + f"numel {numel} ({numel // pack_elems} packs, {world_size} " + f"ranks) with {config} produced {int(wrong.item())} wrong " + "elements" + ) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_pcie_ipc_tiny_payload_every_variant(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _tiny_payload_worker) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_pcie_ipc_all_reduce(world_size: int, dtype: torch.dtype) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip( + f"world_size {world_size} exceeds the {torch.cuda.device_count()} available GPUs" + ) + multi_process_parallel(world_size, _correctness_worker, (dtype,)) + + +def _second_stream_worker(world_size: int, rank: int, port: int) -> None: + """A workspace is bound to one stream; the second one must be rejected. + + The epoch and arrival counters that drive the scratch double buffer are + advanced by the kernels and only well defined if the calls sharing the + workspace are totally ordered. Two streams do not give that. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + hidden = {2: 2048, 4: 4096, 8: 6144}[world_size] + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=hidden * 8, dtype=torch.bfloat16 + ) + inp = torch.randn(8, hidden, dtype=torch.bfloat16, device=device) + + ws.all_reduce(inp) # binds to the current stream + ws.all_reduce(inp) # same stream, still fine + + other = torch.cuda.Stream(device=device) + with ( + torch.cuda.stream(other), + pytest.raises(RuntimeError, match="already bound to"), + ): + ws.all_reduce(inp) + torch.cuda.synchronize(device) + + # And the escape hatch works: once the caller has ordered the two + # streams, rebind_stream() lets the move through and the results stay + # correct. Rejecting the second stream would be no use if the only way + # past it were to build another workspace. + ref = inp.clone() + dist.all_reduce(ref, group=group) + ws.rebind_stream() + with torch.cuda.stream(other): + torch.testing.assert_close(ws.all_reduce(inp), ref, rtol=0, atol=0) + torch.testing.assert_close(ws.all_reduce(inp), ref, rtol=0, atol=0) + torch.cuda.synchronize(device) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [2]) +def test_pcie_ipc_second_stream_rejected(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _second_stream_worker) + + +@pytest.mark.parametrize("world_size", [2]) +def test_pcie_ipc_unsupported_shapes(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _unsupported_shape_worker) + + +_TUNE_HIDDEN = 6144 +_TUNE_BATCHES = (1, 2, 4) + + +def _tune_cache_path(tmpdir: str) -> str: + return os.path.join(tmpdir, "pcie_ipc_tune.json") + + +def _tune_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """Tune, then check the group agrees, the answer is right, and it persists. + + Autotuning a collective is only safe if every rank ends up on the same + kernel. Nothing in the kernels checks that -- they spin without a timeout -- + so the agreement is asserted here rather than assumed. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + path = _tune_cache_path(tmpdir) + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * max(_TUNE_BATCHES), + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=_TUNE_BATCHES, + ) + before = {b: ws.supports(_tune_input(b, device)) for b in _TUNE_BATCHES} + ws.tune([_TUNE_HIDDEN], cache=path, warmup=2, repeat=5) + + resolved = {} + for batch in _TUNE_BATCHES: + inp = _tune_input(batch, device) + # Tuning must not change which shapes are claimed: admission needs a + # comparison against the caller's fallback, and that fallback is not + # one of the candidates. + assert ws.supports(inp) == before[batch], batch + config = ws.tuned_launch_config(inp) + assert config is not None + resolved[batch] = repr(config) + + ref = inp.clone() + dist.all_reduce(ref, group=group) + wrong = torch.tensor( + [int((ws.all_reduce(inp) != ref).sum().item())], device=device + ) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, ( + f"batch {batch} is wrong after tuning with {config}" + ) + + gathered = [None] * world_size + dist.all_gather_object(gathered, resolved, group=group) + assert all(g == gathered[0] for g in gathered), ( + f"ranks resolved different configurations: {gathered}" + ) + if rank == 0: + assert os.path.isfile(path), "rank 0 must persist the tuned cache" + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _tune_input(batch: int, device: torch.device) -> torch.Tensor: + return torch.randint( + 0, 16, (batch, _TUNE_HIDDEN), dtype=torch.int32, device=device + ).to(torch.bfloat16) + + +def _tune_reuse_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """A second process must reuse the persisted cache without re-tuning. + + Loading has to be eager and identical on every rank. A rank that picked up + the file later than its peers -- or dropped entries the others kept -- + would run a different kernel, and the group hangs rather than errors. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + from flashinfer.autotuner import AutoTuner + + AutoTuner.get().clear_cache() + AutoTuner.get().load_configs(_tune_cache_path(tmpdir)) + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * max(_TUNE_BATCHES), + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=_TUNE_BATCHES, + ) + assert not AutoTuner.get().is_tuning_mode, ( + "reuse must not enter tuning mode; a collective sweep here would be " + "a several-second stall in a serving process" + ) + resolved, differs = {}, 0 + for batch in _TUNE_BATCHES: + inp = _tune_input(batch, device) + config = ws.tuned_launch_config(inp) + resolved[batch] = repr(config) + if config != ws.launch_config(inp): + differs += 1 + ref = inp.clone() + dist.all_reduce(ref, group=group) + wrong = torch.tensor( + [int((ws.all_reduce(inp) != ref).sum().item())], device=device + ) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, batch + gathered = [None] * world_size + dist.all_gather_object(gathered, resolved, group=group) + assert all(g == gathered[0] for g in gathered), gathered + # If the reloaded cache never disagreed with the seed, the test would + # pass while proving nothing about persistence. + assert differs > 0, ( + "the persisted cache resolved to the seed configuration everywhere, " + "so this test did not exercise the cache" + ) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _tune_cardinality_worker( + world_size: int, rank: int, port: int, tmpdir: str +) -> None: + """Every rank must issue the identical number of launches while tuning. + + The whole search is a chain of unsynchronised collectives. One rank issuing + a different number of them -- because it screened a different candidate set, + or returned early -- leaves the rest spinning with no timeout. Counting is + the only way to observe it from outside. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * 2, + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=(1, 2), + ) + launches = 0 + original = ws._launch + + def counting(*args, **kwargs): + nonlocal launches + launches += 1 + return original(*args, **kwargs) + + ws._launch = counting + try: + ws.tune([_TUNE_HIDDEN], cache=_tune_cache_path(tmpdir), warmup=2, repeat=5) + finally: + ws._launch = original + + counts = [None] * world_size + dist.all_gather_object(counts, launches, group=group) + assert len(set(counts)) == 1, ( + f"ranks issued different numbers of launches while tuning: {counts}" + ) + assert counts[0] > 0 + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +def _tune_gate_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """A candidate that computes the wrong answer must not survive screening. + + The autotuner ranks purely on time, and this protocol's failure mode is a + sentinel poll returning stale data -- wrong *and* fast. The corruption is + injected in Python rather than by rebuilding a broken kernel: the point is + to test the gate, and a genuinely broken build can spin instead of finish. + """ + ws = None + group = None + tune_group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + from flashinfer.autotuner import set_autotune_process_group + + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * 2, + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=(1, 2), + ) + # Screening only runs when a real search is safe, so install the group + # the search would have. + tune_group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + set_autotune_process_group(tune_group) + from flashinfer.comm.pcie_ipc_tuning import candidate_tactics, tactic_to_config + + poisoned = tactic_to_config(candidate_tactics(world_size, ws.max_blocks)[3]) + original = ws._launch + + def sabotage(inp, out, config, enable_pdl=False): + original(inp, out, config, enable_pdl) + if config == poisoned: + out.add_(1) + + ws._launch = sabotage + try: + inp = _tune_input(2, device) + survivors = ws._runner.get_valid_tactics([inp], None) + finally: + ws._launch = original + set_autotune_process_group(None) + + from flashinfer.comm.pcie_ipc_tuning import config_to_tactic + + assert config_to_tactic(poisoned) not in survivors, ( + f"{poisoned} computed the wrong answer but survived screening" + ) + assert len(survivors) > 1, "screening must not reject everything" + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if tune_group is not None: + dist.destroy_process_group(tune_group) + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_tune_cardinality(world_size: int, tmp_path) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel( + world_size, _tune_cardinality_worker, (str(tmp_path),), timeout_s=900 + ) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_tune_gate_excludes_wrong_candidates( + world_size: int, tmp_path +) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel( + world_size, _tune_gate_worker, (str(tmp_path),), timeout_s=600 + ) + + +def _untuned_warning_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """The untuned warning fires once, and only when the machine is untuned. + + Three ways to get this wrong, all user-visible. Warning per call turns a + serving loop into a log flood. Warning while ``autotune(True)`` is open + advises the caller to tune in the middle of tuning -- that path skips the + hot cache, so every call reaches the cold path. And warning after + :meth:`tune` tells them to redo work they just did. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + from flashinfer.autotuner import AutoTuner, autotune + + AutoTuner.get().clear_cache() + path = os.path.join(tmpdir, f"absent_ws{world_size}.json") + assert not os.path.exists(path) + + def _build(): + return comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * max(_TUNE_BATCHES), + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=_TUNE_BATCHES, + tune_cache=path, + ) + + def _untuned_warnings(fn): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fn() + return [ + w + for w in caught + if "PCIe IPC all-reduce" in str(w.message) + and issubclass(w.category, UserWarning) + ] + + ws = _build() + first = _untuned_warnings(lambda: ws.all_reduce(_tune_input(1, device))) + assert len(first) == 1, first + # A different shape, so this is another cold-path call rather than a hot + # cache hit: the flag has to be what stops it, not the shape cache. + again = _untuned_warnings(lambda: ws.all_reduce(_tune_input(2, device))) + assert not again, again + + # A fresh workspace on the same untuned machine must not warn while it + # is the one doing the tuning. + ws.destroy() + ws = _build() + + def _tune_in_process(): + with autotune(True, tuning_buckets=_TUNE_BATCHES, round_up=False): + for batch in _TUNE_BATCHES: + ws.tuned_launch_config(_tune_input(batch, device)) + + during = _untuned_warnings(_tune_in_process) + assert not during, during + + # And silent afterwards, including on a shape tuning did not cover. + ws.tune([_TUNE_HIDDEN], cache=path, warmup=2, repeat=5) + after = _untuned_warnings(lambda: ws.all_reduce(_tune_input(4, device))) + assert not after, after + + # max_numel is part of the cache key, so a workspace differing only in + # it reads the same file and misses every entry. + ws.destroy() + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=2 * _TUNE_HIDDEN * max(_TUNE_BATCHES), + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=_TUNE_BATCHES, + tune_cache=path, + ) + mismatch = _untuned_warnings(lambda: ws.all_reduce(_tune_input(1, device))) + assert len(mismatch) == 1, mismatch + assert "holds no entry for this workspace" in str(mismatch[0].message) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_untuned_warning(world_size: int, tmp_path) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _untuned_warning_worker, (str(tmp_path),)) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_tune_end_to_end(world_size: int, tmp_path) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + tmpdir = str(tmp_path) + multi_process_parallel(world_size, _tune_worker, (tmpdir,), timeout_s=1200) + multi_process_parallel(world_size, _tune_reuse_worker, (tmpdir,), timeout_s=600) + + +def _standard_idiom_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """The library's usual ``with autotune(True)`` idiom must work here too. + + ``tune()`` is a convenience, not the only entry point -- a caller who knows + FlashInfer should not have to learn a second way to tune one operator. + """ + ws = None + group = None + tune_group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + from flashinfer.autotuner import ( + AutoTuner, + autotune, + set_autotune_process_group, + ) + + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * 2, + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=(1, 2), + ) + path = os.path.join(tmpdir, "idiom.json") + tune_group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + set_autotune_process_group(tune_group) + try: + with autotune(True, cache=path): + for _ in range(2): + for batch in (1, 2): + ws.all_reduce(_tune_input(batch, device)) + finally: + set_autotune_process_group(None) + + resolved = {} + for batch in (1, 2): + inp = _tune_input(batch, device) + ref = inp.clone() + dist.all_reduce(ref, group=group) + wrong = torch.tensor( + [int((ws.all_reduce(inp) != ref).sum().item())], device=device + ) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, batch + resolved[batch] = repr(ws.tuned_launch_config(inp)) + gathered = [None] * world_size + dist.all_gather_object(gathered, resolved, group=group) + assert all(g == gathered[0] for g in gathered), gathered + assert AutoTuner.get().stats.tuned_op_successful_configs.get( + comm.PCIE_IPC_CUSTOM_OP + ), "the standard idiom must actually have profiled something" + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if tune_group is not None: + dist.destroy_process_group(tune_group) + if group is not None: + dist.destroy_process_group(group) + + +def _stray_tuning_worker(world_size: int, rank: int, port: int, tmpdir: str) -> None: + """A tuning session that forgot the reduction group must degrade, not hang. + + Tuning mode is process-global, so a caller tuning its GEMMs sweeps this + operator too. Without a reduction over the candidate timings every rank + would argmin independently and they would not agree on a kernel -- which + this protocol does not survive. The search has to notice and decline. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + from flashinfer.autotuner import autotune, get_autotune_process_group + + assert get_autotune_process_group() is None + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_TUNE_HIDDEN * 2, + dtype=torch.bfloat16, + profile="rootcplx", + tune_batches=(1, 2), + ) + with pytest.warns(RuntimeWarning, match="no matching"), autotune(True): + for batch in (1, 2): + inp = _tune_input(batch, device) + ref = inp.clone() + dist.all_reduce(ref, group=group) + wrong = torch.tensor( + [int((ws.all_reduce(inp) != ref).sum().item())], device=device + ) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, batch + # The seed's answer, unchanged: nothing was searched. + assert ws.tuned_launch_config(inp) == ws.launch_config(inp) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_tune_via_standard_idiom(world_size: int, tmp_path) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel( + world_size, _standard_idiom_worker, (str(tmp_path),), timeout_s=900 + ) + + +@pytest.mark.parametrize("world_size", [8]) +def test_pcie_ipc_stray_tuning_session_degrades(world_size: int, tmp_path) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel( + world_size, _stray_tuning_worker, (str(tmp_path),), timeout_s=300 + ) + + +def _graph_capture_worker(world_size: int, rank: int, port: int) -> None: + """A shape must be prepared before it can be captured, and say so if not. + + Resolving a launch configuration makes the group agree on it with a + reduction whose verdict is read back on the host, and neither the + collective nor the readback is legal inside a capture. Both halves are + pinned here: the unprepared shape must fail with an error that names the + remedy, and the prepared one must capture, replay, and still be correct. + """ + ws = None + group = None + try: + _init_process_group(world_size, rank, port) + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + hidden, batch = 4096, 4 + ws = comm.PcieIpcAllReduceWorkspace( + group=group, max_numel=batch * hidden, dtype=dtype + ) + inp = torch.randint( + 0, 16, (batch, hidden), dtype=torch.int32, device=device + ).to(dtype) + out = torch.empty_like(inp) + + # The op has to be usable at all before the capture claims mean anything. + if not ws.supports(inp): + return + reference = inp.clone() + dist.all_reduce(reference, group=group) + + # Unprepared: the guard fires before any collective is issued, so every + # rank raises at the same point and the group stays in step. + torch.cuda.synchronize(device) + dist.barrier(group=group) + unprepared = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(unprepared): + ws.all_reduce(inp, out=out) + raised = "" + except RuntimeError as exc: + raised = str(exc) + assert "prepare()" in raised, ( + "capturing an unresolved shape should name prepare() as the " + f"remedy, got: {raised or 'no error at all'}" + ) + + # Prepared: capture, replay, and check the result really is the sum. + ws.prepare([(batch, hidden)], dtype=dtype) + torch.cuda.synchronize(device) + dist.barrier(group=group) + prepared = torch.cuda.CUDAGraph() + with torch.cuda.graph(prepared): + ws.all_reduce(inp, out=out) + out.zero_() + prepared.replay() + torch.cuda.synchronize(device) + + wrong = torch.tensor([int((out != reference).sum().item())], device=device) + dist.all_reduce(wrong, op=dist.ReduceOp.MAX, group=group) + assert int(wrong.item()) == 0, ( + f"the replayed graph produced {int(wrong.item())} wrong elements" + ) + dist.barrier(group=group) + finally: + if ws is not None: + ws.destroy() + if group is not None: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_pcie_ipc_prepare_makes_a_shape_capturable(world_size: int) -> None: + if world_size > torch.cuda.device_count(): + pytest.skip("not enough GPUs") + multi_process_parallel(world_size, _graph_capture_worker, timeout_s=300) diff --git a/tests/comm/test_pcie_ipc_cross_island_race.py b/tests/comm/test_pcie_ipc_cross_island_race.py new file mode 100644 index 00000000000..e43c3edaf24 --- /dev/null +++ b/tests/comm/test_pcie_ipc_cross_island_race.py @@ -0,0 +1,212 @@ +""" +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. + +Deterministic regression for the TP8 cross-island scratch race. + +The block scratch is double buffered so that a fast island cannot overwrite a +cross slot the slow island has not read yet. That property cannot be tested by +running the collective harder: the window only opens when one island stalls +*between* the owner-pair rendezvous and its cross read, which a host-side delay +cannot produce because the pair barrier releases both islands together. + +So the kernel carries a debug-only stall, and this file pins the behaviour from +both sides: + + * with the stall and the double buffer -> correct + * with the stall and the double buffer disabled -> wrong + +The negative control is the point. Without it, a passing test says nothing +about whether the fix is load-bearing. + +Opt-in: needs 8 GPUs and builds two extra JIT modules (~1 min each), so it is +skipped unless FLASHINFER_TEST_PCIE_IPC_RACE=1. +""" + +import os + +import pytest +import torch +import torch.distributed as dist + +from tests.comm.test_pcie_ipc_all_reduce import ( + _init_process_group, + multi_process_parallel, +) + +pytestmark = pytest.mark.skipif( + os.environ.get("FLASHINFER_TEST_PCIE_IPC_RACE") != "1", + reason="opt-in: needs 8 GPUs and builds two extra JIT modules", +) + +_HIDDEN = 6144 +# 100 us reproduced at batch 16 (one collective is ~63 us there); 200 us leaves +# margin on a busier machine without making the test slow. +_STALL_NS = 200_000 +_ITERS = 40 + + +def _mismatches(world_size: int, rank: int, stall_ns: int, no_block_epoch: int) -> int: + """Run the race sequence and return the group-wide mismatch count.""" + import flashinfer.comm as comm + from flashinfer.comm import pcie_ipc_ar + from flashinfer.jit.comm import gen_pcie_ipc_comm_debug_module + + # The workspace resolves its module through this module-global, so swapping + # it here is enough to put the instrumented build under the same wrapper + # the shipping path uses. + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + original = pcie_ipc_ar.gen_pcie_ipc_comm_module + pcie_ipc_ar.gen_pcie_ipc_comm_module = lambda: gen_pcie_ipc_comm_debug_module( + stall_ns, 0, no_block_epoch + ) + try: + group = dist.group.WORLD + device = torch.device(f"cuda:{rank}") + batches = [16, 16, 96, 16] # exercises ring blocks 1 -> 1 -> 2 -> 1 + # Pin the profile. The batches above are chosen against the rootcplx + # table, where they all select the ring kernel -- the one the stall is + # instrumented in. On a switch-paired machine the same batches select + # the block kernel instead, and the test would still pass while + # exercising something other than what it claims to. + ws = comm.PcieIpcAllReduceWorkspace( + group=group, + max_numel=_HIDDEN * max(batches), + dtype=torch.bfloat16, + profile="rootcplx", + ) + try: + # Distinct payloads per call are essential. Reusing one input makes + # every call produce the same island partial sums, so the racing + # write stores a bit-identical value over the one it clobbers and + # nothing is observable however wide the window is. + inputs = [ + torch.randint(0, 16, (b, _HIDDEN), dtype=torch.int32, device=device).to( + torch.bfloat16 + ) + + i + for i, b in enumerate(batches) + ] + refs = [] + for x in inputs: + r = x.clone() + dist.all_reduce(r, group=group) + refs.append(r) + dist.barrier(group=group) + + bad = 0 + for _ in range(_ITERS): + for x, ref in zip(inputs, refs, strict=True): + bad += int((ws.all_reduce(x) != ref).sum().item()) + finally: + ws.destroy() + + total = torch.tensor([bad], device=device) + dist.all_reduce(total, group=group) + return int(total.item()) + finally: + pcie_ipc_ar.gen_pcie_ipc_comm_module = original + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + + +def _fixed_worker(world_size: int, rank: int, port: int) -> None: + try: + _init_process_group(world_size, rank, port) + bad = _mismatches(world_size, rank, _STALL_NS, no_block_epoch=0) + assert bad == 0, ( + f"{bad} mismatched elements with the double buffer enabled: the " + "cross-island scratch is being reused too early" + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _negative_control_worker(world_size: int, rank: int, port: int) -> None: + try: + _init_process_group(world_size, rank, port) + bad = _mismatches(world_size, rank, _STALL_NS, no_block_epoch=1) + assert bad > 0, ( + "no mismatch with the double buffer disabled -- this test is not " + "exercising the race, so its passing sibling proves nothing. " + "Check that the stall is still between the pair rendezvous and the " + "cross read, and that the payloads differ between calls." + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _grid_change_worker(world_size: int, rank: int, port: int) -> None: + """The shipping sequence from test_pcie_ipc_grid_change, on the old scheme. + + Built with per-block epoch parity, so growing the grid puts a + first-appearance block on the half the previous call used. Here that is not + a corrupted number but a deadlock: the mis-addressed write lands on the + address the victim polls next, so the victim waits for a sentinel nobody + will clear. + """ + from tests.comm.test_pcie_ipc_all_reduce import ( + _GRID_CHANGE_BATCHES, + _grid_change_worker as shipping_worker, + ) + + from flashinfer.comm import pcie_ipc_ar + from flashinfer.jit.comm import gen_pcie_ipc_comm_debug_module + + assert _GRID_CHANGE_BATCHES # the sequence under test is the shipping one + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + original = pcie_ipc_ar.gen_pcie_ipc_comm_module + pcie_ipc_ar.gen_pcie_ipc_comm_module = lambda: gen_pcie_ipc_comm_debug_module( + 0, 0, 0, 1 + ) + try: + shipping_worker(world_size, rank, port) + finally: + pcie_ipc_ar.gen_pcie_ipc_comm_module = original + pcie_ipc_ar.get_pcie_ipc_comm_module.cache_clear() + + +def _requires_8_gpus() -> None: + if torch.cuda.device_count() < 8: + pytest.skip("the cross-island race needs both 4-GPU islands") + + +def test_cross_island_scratch_is_double_buffered() -> None: + _requires_8_gpus() + multi_process_parallel(8, _fixed_worker) + + +def test_cross_island_race_reproduces_without_double_buffer() -> None: + """Negative control: the fix must be load-bearing.""" + _requires_8_gpus() + multi_process_parallel(8, _negative_control_worker) + + +def test_grid_change_deadlocks_with_per_block_epoch() -> None: + """Negative control for the call-level epoch: the old scheme must deadlock. + + Asserting on the deadlock rather than on a mismatch count is not a + compromise, it is the actual signal. In these kernels a mis-addressed write + lands on the address the victim polls next, so corruption and deadlock are + the same event -- there is no interleaving that yields a wrong number + without also wedging the spin. + + The timeout is short because the hang is immediate: it happens on the third + call, within a second of the workspace being built. + """ + if torch.cuda.device_count() < 4: + pytest.skip("the grid-change sequence is tuned against the 4-rank table") + with pytest.raises(AssertionError, match="did not finish"): + multi_process_parallel(4, _grid_change_worker, timeout_s=90) diff --git a/tests/comm/test_pcie_ipc_policy.py b/tests/comm/test_pcie_ipc_policy.py new file mode 100644 index 00000000000..59dd0f00699 --- /dev/null +++ b/tests/comm/test_pcie_ipc_policy.py @@ -0,0 +1,522 @@ +""" +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. + +Properties of the PCIe IPC launch policy. + +The policy is a pure function, and everything else about this collective rests +on that: every rank derives its own launch configuration with no runtime +agreement, so a policy that answered differently on two ranks -- or that +answered with a configuration the kernel rejects -- would hang the group rather +than fail. + +Two layers with different standing, tested differently. *Admission* is a +capability claim and is pinned exactly: it decides which shapes reach the +kernels at all, and both of its answers are load-bearing -- a false yes reaches +a hard check mid-collective, a false no silently routes a supported shape to +another backend. The *seed* is a default rather than a measurement, so only its +shape is asserted here; its constants belong to whatever machine measured them +and are re-measured by ``PcieIpcAllReduceWorkspace.tune``. + +Nothing here needs a GPU, and nothing here may start to: no CUDA, no process +group. +""" + +import importlib.util +import inspect +import pathlib + +import pytest + +from flashinfer.comm import pcie_ipc_tuning as tuning +from flashinfer.comm.pcie_ipc_policy import ( + MAX_BLOCKS, + IpcLaunchConfig, + IpcVariant, + _admits, + _is_launchable, + _seed, + get_pcie_ipc_launch_config, +) + +# The launcher accepts 2-byte dtypes only (bfloat16, float16), so this is the +# only element size a caller can reach. +_ELEM_SIZE = 2 +_PACK_ELEMS = 16 // _ELEM_SIZE + +_WORLD_SIZES = (2, 4, 8) +_HIDDENS = (1024, 2048, 4096, 6144, 8192) +_BATCHES = (1, 2, 3, 4, 8, 16, 32, 64, 128, 256) +_NUMELS = tuple(sorted({b * h for b in _BATCHES for h in _HIDDENS})) + +# Eight packs up to a few million elements, doubling. Wide enough to contain any +# plausible crossover from either direction without naming where it sits. +_LADDER = tuple(64 << k for k in range(17)) + + +def _config(world_size, numel, max_blocks=MAX_BLOCKS): + return get_pcie_ipc_launch_config(world_size, numel, _ELEM_SIZE, max_blocks) + + +@pytest.mark.parametrize("world_size", _WORLD_SIZES) +def test_every_returned_config_is_launchable(world_size: int) -> None: + """A config the kernel would reject must never leave the policy. + + The C++ side hard-checks these. Reaching them means one rank raises while + its peers are already spinning in the collective. + """ + admitted = 0 + for numel in _NUMELS: + for max_blocks in (MAX_BLOCKS, 8): + config = _config(world_size, numel, max_blocks) + if config is None: + continue + admitted += 1 + assert _is_launchable(world_size, config, max_blocks), ( + f"ws={world_size} numel={numel} max_blocks={max_blocks} " + f"-> {config}, which the kernel rejects" + ) + assert 0 < config.blocks <= max_blocks + assert world_size <= config.threads <= 1024 + assert admitted, "the sweep admits nothing, so it checks nothing" + + +def test_the_policy_is_a_pure_function() -> None: + """Same arguments, same answer -- no hidden state, no rank-local input.""" + keys = [(w, n) for w in _WORLD_SIZES for n in _NUMELS] + first = {k: _config(*k) for k in keys} + assert any(v is not None for v in first.values()) + # The answers are memoised, so asking again would only re-read the cache. + # Drop it and recompute in the opposite order, which re-runs the function + # and would also expose an answer that depended on call order. + get_pcie_ipc_launch_config.cache_clear() + assert {k: _config(*k) for k in reversed(keys)} == first + + +def test_flat_staged_is_never_selected_outside_world_size_eight() -> None: + """It would name the same kernel as ``STAGED`` at 4 ranks, and none at 2.""" + for world_size in (2, 4): + for numel in _NUMELS: + # The seed, not the gated return value: the gate below already + # refuses FLAT_STAGED here, so it would answer for the seed and the + # loop could not fail. + seed = _seed(world_size, numel, _ELEM_SIZE, MAX_BLOCKS) + assert seed.variant is not IpcVariant.FLAT_STAGED + # And the launchability check refuses it even if the seed ever returned it. + bad = IpcLaunchConfig(1, 128, IpcVariant.FLAT_STAGED) + assert not _is_launchable(4, bad, MAX_BLOCKS) + assert not _is_launchable(2, bad, MAX_BLOCKS) + assert _is_launchable(8, bad, MAX_BLOCKS) + + +def test_world_size_eight_staged_always_gets_a_multiple_of_four_blocks() -> None: + """The block-partitioned TP8 kernel derives its chunk from ``blockIdx.x & 3``. + + The seed is not the only source of configurations for that kernel and need + not reach it at all, so the tuner's candidates are checked too -- they share + the gate, and the gate is what has to hold. + """ + for numel in _NUMELS: + # Same reason as the FLAT_STAGED case: read the seed, since the gate + # would otherwise turn a bad block count into None and a green test. + seed = _seed(8, numel, _ELEM_SIZE, MAX_BLOCKS) + if seed.variant is IpcVariant.STAGED: + assert seed.blocks % 4 == 0, f"numel={numel} -> {seed}" + + staged = [ + config + for config in map(tuning.tactic_to_config, tuning.candidate_tactics(8)) + if config.variant is IpcVariant.STAGED + ] + assert staged, "the tuner offers this kernel nothing to run" + for config in staged: + assert config.blocks % 4 == 0, config + + assert not _is_launchable(8, IpcLaunchConfig(2, 256, IpcVariant.STAGED), MAX_BLOCKS) + assert _is_launchable(8, IpcLaunchConfig(4, 256, IpcVariant.STAGED), MAX_BLOCKS) + + +@pytest.mark.parametrize("world_size", _WORLD_SIZES) +def test_admission_floor_is_one_pack_per_rank(world_size: int) -> None: + """Below one 16-byte pack per rank the reduce-scatter split degenerates. + + Refused rather than served, because the two ownership formulas in the + kernels stop agreeing there and a payload that small belongs on another + backend anyway. + """ + floor = _PACK_ELEMS * world_size + assert _config(world_size, floor) is not None + # A whole pack below the floor, so this isolates the per-rank rule from the + # whole-pack rule tested separately. + assert _config(world_size, floor - _PACK_ELEMS) is None + assert _config(world_size, _PACK_ELEMS) is None + + +def test_the_admission_floor_scales_with_the_world_size() -> None: + """One shape, admitted or not depending only on how many ranks share it.""" + numel = _PACK_ELEMS * 4 + assert _config(2, numel) is not None + assert _config(4, numel) is not None + assert _config(8, numel) is None + + +@pytest.mark.parametrize("world_size", _WORLD_SIZES) +def test_numel_that_is_not_whole_packs_is_refused(world_size: int) -> None: + """The kernels address the payload in 16-byte packs; the launcher agrees.""" + whole = _PACK_ELEMS * 64 + assert _config(world_size, whole) is not None + for remainder in (1, _PACK_ELEMS // 2, _PACK_ELEMS - 1): + assert _config(world_size, whole + remainder) is None, remainder + + +@pytest.mark.parametrize("world_size", (0, 1, 3, 5, 6, 7, 9, 16)) +def test_world_sizes_the_kernels_do_not_implement_are_refused(world_size: int) -> None: + for numel in (_PACK_ELEMS * 128, 6144, 65536): + assert _config(world_size, numel) is None, numel + + +@pytest.mark.parametrize( + "world_size,hidden", [(4, 2048), (4, 8192), (2, 4096), (8, 1024), (8, 2048)] +) +def test_admission_does_not_depend_on_the_hidden_size( + world_size: int, hidden: int +) -> None: + """Hidden size is not a term in a capability question. + + These pairs are the ones a per-hidden restriction singles out first, and the + kernels run all of them; refusing one would route a supported shape to + another backend for good, since a caller reads ``None`` as "use NCCL". + """ + for batch in (1, 2, 7, 64, 256): + config = _config(world_size, batch * hidden) + assert config is not None, f"ws={world_size} hidden={hidden} batch={batch}" + assert _is_launchable(world_size, config, MAX_BLOCKS) + + +def test_the_seed_crosses_to_a_staged_kernel_as_the_payload_grows() -> None: + """The one crossover that ports between machines, asserted as a shape. + + One-shot moves ``(N-1)*P`` bytes per rank in a single round trip; staging + moves ``2*(N-1)*P/N`` and adds barriers, so it wins once the payload is + large. Where it wins is a property of the machine and is measured by + ``tune``; *that* it wins, once and without coming back, is not, and is what + is pinned here. + """ + for world_size in (4, 8): + variants = [_config(world_size, numel).variant for numel in _LADDER] + assert variants[0] is not IpcVariant.STAGED_RING, world_size + assert variants[-1] is IpcVariant.STAGED_RING, world_size + staged = [v is IpcVariant.STAGED_RING for v in variants] + assert staged == sorted(staged), ( + f"ws={world_size} leaves the staged kernel as the payload grows: " + f"{list(zip(_LADDER, variants, strict=True))}" + ) + + # Two ranks stage exactly the bytes they would have pushed (2P/N is P at + # N == 2), so there is no crossover to place and no second branch. + assert len({_config(2, numel).variant for numel in _LADDER}) == 1 + + +def test_the_answer_keys_on_the_element_count_not_on_its_factorisation() -> None: + """A token count does not port between machines; a byte count does. + + A policy keyed on ``(hidden, batch)`` answers one payload two ways depending + on which shape produced it. The signature is what forecloses that: with + ``numel`` the only shape argument, no factorisation can reach the function. + """ + params = list(inspect.signature(get_pcie_ipc_launch_config).parameters) + assert params[:3] == ["world_size", "numel", "elem_size"] + assert not {"batch", "hidden", "profile"} & set(params), params + + +@pytest.mark.parametrize("max_blocks", (1, 2, 3, 4, 7, 16, MAX_BLOCKS)) +def test_a_small_max_blocks_is_never_exceeded(max_blocks: int) -> None: + """The workspace is sized for its own ``max_blocks``, not for the default. + + A configuration over that budget indexes scratch that was never allocated, + which the launcher rejects on the rank that asked and nowhere else. + """ + for world_size in _WORLD_SIZES: + admitted = 0 + for numel in _NUMELS: + config = _config(world_size, numel, max_blocks) + # A tighter max_blocks bounds the grid, never the supported set. + assert (config is not None) == _admits(world_size, numel, _ELEM_SIZE) + if config is None: + continue + admitted += 1 + assert 0 < config.blocks <= max_blocks, (world_size, numel, config) + assert _is_launchable(world_size, config, max_blocks) + assert admitted, f"ws={world_size} max_blocks={max_blocks} checks nothing" + + +def test_every_dispatchable_variant_is_reachable_from_the_tuner() -> None: + """A variant the dispatch can launch but nothing can select is invisible. + + This is the direction ``test_every_returned_config_is_launchable`` does not + cover. A kernel that nothing selects is either dead code or an unexplored + region of the launch space, and the two are indistinguishable from outside + -- which is how the flat-staged kernel stayed unreachable for a whole tuning + round. The seed deliberately reaches only some of them, so the tuner's + candidate list is where this has to hold. + """ + expected = { + 2: {IpcVariant.UNSTAGED, IpcVariant.STAGED}, + 4: {IpcVariant.UNSTAGED, IpcVariant.STAGED, IpcVariant.STAGED_RING}, + 8: { + IpcVariant.UNSTAGED, + IpcVariant.STAGED, + IpcVariant.STAGED_RING, + IpcVariant.FLAT_STAGED, + }, + } + seen = { + world_size: { + tuning.tactic_to_config(t).variant + for t in tuning.candidate_tactics(world_size) + } + for world_size in _WORLD_SIZES + } + assert seen == expected + + +def _load_benchmark_module(): + """Import the benchmark as a module so its pure helpers can be tested.""" + path = ( + pathlib.Path(__file__).resolve().parents[2] + / "benchmarks" + / "comm" + / "bench_pcie_ipc_all_reduce.py" + ) + spec = importlib.util.spec_from_file_location("_bench_pcie_ipc", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _one_config_per_variant(world_size): + """One launchable configuration per variant the dispatch reaches.""" + out = {} + for tactic in tuning.candidate_tactics(world_size): + config = tuning.tactic_to_config(tactic) + out.setdefault(config.variant, config) + return out + + +def _expected_switch(world_size, config): + """The protocol this kernel actually used to run. + + TP2, TP4, the TP8 pack kernel and the flat-staged kernel double-buffered by + per-block parity. The two topology-staged TP8 kernels had no epoch at all. + """ + topo_staged_tp8 = world_size == 8 and config.variant in ( + IpcVariant.STAGED, + IpcVariant.STAGED_RING, + ) + return "no-block-epoch" if topo_staged_tp8 else "per-block-epoch" + + +def test_protocol_ab_picks_the_right_history_per_kernel() -> None: + """The A/B baseline must match what each kernel actually used to run. + + Comparing against the wrong one measures a protocol that never shipped, + which is not a performance result. Checked over every kernel the dispatch + reaches rather than only the ones some shape currently lands on, since which + kernel a shape gets is a default that is expected to move. + """ + bench = _load_benchmark_module() + pick = bench._historical_switch + + reached = set() + for world_size in _WORLD_SIZES: + for variant, config in _one_config_per_variant(world_size).items(): + want = _expected_switch(world_size, config) + reached.add(want) + assert pick(world_size, config) == want, ( + f"ws={world_size} {variant.name} -> {pick(world_size, config)}, " + f"expected {want}" + ) + assert reached == {"per-block-epoch", "no-block-epoch"} + + # And the shapes the shipping path actually produces land on the same side. + for world_size, hidden in ((2, 2048), (4, 4096), (8, 6144)): + for batch in (1, 2, 4, 8, 16, 32, 64, 128): + config = _config(world_size, batch * hidden) + if config is None: + continue + assert pick(world_size, config) == _expected_switch(world_size, config), ( + f"ws={world_size} hidden={hidden} batch={batch} -> {config}" + ) + + +def _shape_config(world_size, hidden, batches): + out = {} + for batch in batches: + config = _config(world_size, batch * hidden) + if config is not None: + out[batch] = config + return out + + +def test_protocol_ab_plan_partitions_shapes_by_history() -> None: + """The history split must drive execution, not just labelling. + + Labelling a row SYNTHETIC after the fact does not undo running it, and one + of those runs -- the fixed-half build on a kernel that always + double-buffered -- wedges the sentinel loop, taking the whole benchmark down + before it reports anything. So the plan each leg executes is asserted here, + not only the label it would print. + """ + bench = _load_benchmark_module() + plan_of = bench._protocol_ab_plan + batches = [1, 2, 4, 8, 16, 32, 64, 128] + + # 2 and 4 ranks never needed the fixed-half build, so it must not run there. + for world_size, hidden in ((2, 2048), (4, 4096)): + shapes = _shape_config(world_size, hidden, batches) + assert shapes, f"ws={world_size} hidden={hidden} admits nothing" + plan = plan_of(world_size, shapes, "auto") + assert set(plan) == {"per-block-epoch"} + assert plan["per-block-epoch"] == sorted(shapes) + + # 8 ranks straddles both histories, which is the whole reason auto exists. + shapes = _shape_config(8, 6144, batches) + history = {batch: _expected_switch(8, c) for batch, c in shapes.items()} + assert set(history.values()) == {"per-block-epoch", "no-block-epoch"}, ( + "this sweep no longer covers both histories, so the partition below " + f"would hold for want of anything to separate: {history}" + ) + + plan = plan_of(8, shapes, "auto") + assert set(plan) == {"per-block-epoch", "no-block-epoch"} + per_block = plan["per-block-epoch"] + no_epoch = plan["no-block-epoch"] + assert not set(per_block) & set(no_epoch), ( + "a shape may only run under its own history" + ) + assert sorted(per_block + no_epoch) == sorted(shapes), ( + "every admitted shape must run once" + ) + for switch, leg_batches in plan.items(): + for batch in leg_batches: + assert history[batch] == switch, ( + f"batch {batch} ({shapes[batch]}) ran under {switch}, " + f"but its history is {history[batch]}" + ) + + # An explicit switch is a request for that exact comparison, so it keeps + # every shape and reports the mismatched ones rather than dropping them. + explicit = plan_of(8, shapes, "per-block-epoch") + assert explicit == {"per-block-epoch": sorted(shapes)} + assert plan_of(8, {}, "auto") == {} + + +def test_protocol_ab_legs_execute_only_their_own_shapes() -> None: + """The plan must drive execution, not sit beside it. + + An earlier version of the harness computed the split correctly and then + ignored it, running every shape under every switch and discarding the + mismatched rows afterwards. Discarding is too late: the kernel has already + run, and the fixed-half build on a kernel that always double-buffered wedges + its spin loop rather than returning a bad number. A test of the planner + alone stayed green through all of that, so this one records what each leg is + actually handed. + """ + bench = _load_benchmark_module() + batches = [1, 2, 4, 8, 16, 32, 64, 128] + shapes = _shape_config(8, 6144, batches) + plan = bench._protocol_ab_plan(8, shapes, "auto") + assert len(plan) == 2, f"one leg only; nothing to keep apart: {plan}" + + calls = [] + + def recording_sweep(broken_switch, leg_batches): + calls.append((broken_switch, tuple(leg_batches))) + return {b: ("timing", shapes[b], True, True) for b in leg_batches} + + legs = bench._run_ab_legs(plan, recording_sweep) + + assert set(legs) == set(plan) + for switch, leg_batches in plan.items(): + # A-B-A: broken, shipping, broken -- all three over the same shapes. + assert (switch, tuple(leg_batches)) in calls + assert (None, tuple(leg_batches)) in calls + assert len(legs[switch]) == 3 + for leg in legs[switch]: + assert set(leg) == set(leg_batches) + + # Nothing ran outside its own plan entry. + for broken_switch, ran in calls: + if broken_switch is not None: + assert set(ran) == set(plan[broken_switch]), ( + f"{broken_switch} leg ran {sorted(ran)}, owns {plan[broken_switch]}" + ) + else: + assert any(set(ran) == set(v) for v in plan.values()) + + # Each switch is exercised three times and no more. + for switch, leg_batches in plan.items(): + assert calls.count((switch, tuple(leg_batches))) == 2 # the two A legs + assert calls.count((None, tuple(leg_batches))) == 1 # the shipping leg + + +def test_group_correctness_uses_a_min_reduction() -> None: + """Correctness must be an AND across ranks, which means ReduceOp.MIN. + + The cross-island race this harness rebuilds produces errors confined to one + island, so a rank-local verdict -- or a SUM, or a MAX -- would let a clean + rank 0 publish a cost for a baseline that was wrong on ranks 4-7. That + happened once. The op is checked here rather than on hardware because the + failure is a silent wrong verdict, not a crash, and no GPU is needed to see + which reduction was asked for. + """ + bench = _load_benchmark_module() + calls = [] + + class _FakeTensor: + def __init__(self, value): + self.value = value + + def item(self): + return self.value + + def fake_all_reduce(tensor, op=None, group=None): + calls.append(op) + + real_dist = bench.dist + real_torch = bench.torch + try: + bench.dist = type( + "D", + (), + { + "all_reduce": staticmethod(fake_all_reduce), + "ReduceOp": real_dist.ReduceOp, + }, + ) + bench.torch = type( + "T", + (), + {"tensor": staticmethod(lambda v, **kw: _FakeTensor(v[0])), "int32": None}, + ) + assert bench._group_all(True, None, None) is True + assert bench._group_all(False, None, None) is False + finally: + bench.dist = real_dist + bench.torch = real_torch + + assert calls, "the verdict must be reduced across ranks, not decided locally" + assert all(op is real_dist.ReduceOp.MIN for op in calls), ( + f"correctness must reduce with MIN (logical AND), got {calls}" + ) diff --git a/tests/comm/test_pcie_ipc_tuning.py b/tests/comm/test_pcie_ipc_tuning.py new file mode 100644 index 00000000000..7659c59ce1c --- /dev/null +++ b/tests/comm/test_pcie_ipc_tuning.py @@ -0,0 +1,511 @@ +""" +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. + +Properties of the PCIe IPC autotuning layer that need no GPU. + +Everything here guards a failure that is silent on hardware: a tactic that +cannot be persisted, a cache entry reused for the wrong workspace, a verdict +reduced with the wrong operator. The multi-GPU tests can only observe the +consequences, and one of the consequences is a hang. +""" + +import json + +import pytest +import torch + +from flashinfer.autotuner import _json_to_tactic, _tactic_to_json, make_bucket_mapper +from flashinfer.comm import pcie_ipc_tuning as tuning +from flashinfer.comm.pcie_ipc_ar import PcieIpcAllReduceWorkspace +from flashinfer.comm.pcie_ipc_policy import ( + MAX_BLOCKS, + IpcLaunchConfig, + IpcVariant, + _is_launchable, + get_pcie_ipc_launch_config, +) +from flashinfer.comm.pcie_ipc_topology import PROFILE_ROOTCPLX, PROFILE_SWITCHPAIR + +_WORLD_SIZES = (2, 4, 8) + + +@pytest.mark.parametrize("world_size", _WORLD_SIZES) +def test_candidates_are_pure_and_launchable(world_size: int) -> None: + """Every rank must derive the same candidate list, and none may be rejected. + + The autotuner profiles candidates collectively with no timeout, so a list + that differs by one element between ranks deadlocks on the first timing + reduction. And a candidate the launcher rejects raises on the calling rank + only, leaving its peers spinning. + """ + first = tuning.candidate_tactics(world_size) + assert first == tuning.candidate_tactics(world_size) + assert len(first) == len(set(first)), "candidates must be distinct" + for tactic in first: + config = tuning.tactic_to_config(tactic) + assert _is_launchable(world_size, config, MAX_BLOCKS), tactic + + +def test_candidate_rejections_match_the_documented_rules() -> None: + """The only configurations excluded are the ones the header cannot dispatch.""" + grid = [ + (int(v), b, t) + for v in IpcVariant + for b in tuning.TUNE_BLOCKS + for t in tuning.TUNE_THREADS + ] + + # World size 8: the block-partitioned kernel needs blocks % 4 == 0. + rejected = set(grid) - set(tuning.candidate_tactics(8)) + assert rejected == { + (int(IpcVariant.STAGED), b, t) + for b in tuning.TUNE_BLOCKS + if b % 4 != 0 + for t in tuning.TUNE_THREADS + } + + # World size 4: no FLAT_STAGED, and threads must be at least world_size + # (which every entry in the grid already satisfies). + rejected4 = set(grid) - set(tuning.candidate_tactics(4)) + assert all(t[0] == int(IpcVariant.FLAT_STAGED) for t in rejected4) + + # World size 2: only the two variants that name a TP2 kernel. + for tactic in tuning.candidate_tactics(2): + assert tactic[0] in (int(IpcVariant.UNSTAGED), int(IpcVariant.STAGED)) + + +# Every distinct winner in the tuned caches of the two fabrics this operator +# was developed against, as ``(variant, blocks, threads)``. Measured, not +# derived: they are what a real search picked, per world size. +_MEASURED_WINNERS = { + 2: ( + (0, 8, 64), + (0, 8, 128), + (0, 8, 256), + (0, 8, 512), + (0, 12, 128), + (0, 16, 64), + (0, 16, 256), + (0, 32, 64), + (0, 32, 512), + (0, 64, 64), + (0, 96, 64), + (0, 96, 256), + (0, 128, 64), + (1, 1, 256), + (1, 8, 256), + (1, 12, 256), + (1, 16, 128), + (1, 16, 256), + (1, 16, 1024), + (1, 32, 64), + (1, 64, 64), + (1, 128, 64), + ), + 4: ( + (0, 16, 256), + (0, 32, 512), + (1, 1, 128), + (1, 32, 256), + (1, 32, 1024), + (1, 96, 128), + (1, 96, 1024), + (1, 128, 256), + (2, 1, 1024), + (2, 2, 512), + (2, 2, 1024), + ), + 8: ( + (0, 2, 512), + (0, 8, 256), + (0, 12, 256), + (0, 32, 1024), + (1, 4, 512), + (1, 8, 1024), + (1, 96, 256), + (2, 1, 256), + (2, 1, 512), + (2, 1, 1024), + (2, 2, 1024), + (3, 1, 64), + ), +} + + +def test_the_grid_can_express_every_measured_winner() -> None: + """The grid must contain the configurations searches actually chose. + + A winner the grid cannot name is one the tuner can never pick again, which + caps the operator at whatever the seed happens to guess. Block counts of 12 + and 96 are among them, so a powers-of-two grid would silently have this + property. + """ + assert set(_MEASURED_WINNERS) == set(_WORLD_SIZES), "a world size lost its winners" + for world_size, winners in sorted(_MEASURED_WINNERS.items()): + unreachable = sorted(set(winners) - set(tuning.candidate_tactics(world_size))) + assert not unreachable, ( + f"world size {world_size}: the grid cannot name {unreachable}" + ) + + +@pytest.mark.xfail( + strict=True, + reason="TUNE_BLOCKS has no 3, and the seed asks for 3 ring blocks at every " + "payload between 768 KiB and 1 MiB", +) +def test_the_grid_can_express_the_seed() -> None: + """Tuning must be able to reach the seed, not merely fall back to it. + + Tactic -1 always reproduces the seed, but if the grid cannot name the + seed's own configuration then the tuner can only accept or reject it + wholesale -- it can never search the immediate neighbourhood of the one + configuration every untuned shape runs. + """ + missing = [] + for world_size in _WORLD_SIZES: + grid = set(tuning.candidate_tactics(world_size)) + for hidden in (1024, 2048, 4096, 6144): + for batch in range(1, 129): + # Both supported dtypes are 2 bytes, and the seed keys on + # payload bytes, so one element size covers the space. + config = get_pcie_ipc_launch_config(world_size, batch * hidden, 2) + if config is None: + continue + if tuning.config_to_tactic(config) not in grid: + missing.append((world_size, hidden, batch, config)) + assert not missing, f"grid cannot express these seed configurations: {missing[:5]}" + + +@pytest.mark.parametrize("world_size", _WORLD_SIZES) +def test_tactics_survive_the_persistence_round_trip(world_size: int) -> None: + """A tactic that cannot be serialised fails at save time, after the whole run. + + ``_tactic_to_json`` passes anything that is not a scalar or an iterable + through unchanged, so a dataclass reaches ``json.dump`` and raises there -- + at the end of tuning, with every measurement already discarded. + """ + for tactic in tuning.candidate_tactics(world_size): + encoded = json.loads(json.dumps(_tactic_to_json(tactic))) + assert _json_to_tactic(encoded) == tactic + assert _json_to_tactic(json.loads(json.dumps(_tactic_to_json(-1)))) == -1 + + +def test_tactic_codec_round_trips_and_rejects_nonsense() -> None: + for world_size in _WORLD_SIZES: + for tactic in tuning.candidate_tactics(world_size): + assert tuning.config_to_tactic(tuning.tactic_to_config(tactic)) == tactic + with pytest.raises(ValueError): + tuning.tactic_to_config((0, 1)) + with pytest.raises(ValueError): + tuning.tactic_to_config((99, 1, 128)) + + +def test_resolve_falls_back_to_the_seed() -> None: + """Every way a tactic can be unusable ends at the seed, not at an exception. + + The autotuner does not check that a cached tactic can implement the shape + it is reused for, so a stale entry has to be caught here. Raising instead + would take down one rank mid-collective and leave the rest spinning. + """ + seed = IpcLaunchConfig(32, 128, IpcVariant.UNSTAGED) + resolve = tuning.resolve_tuned_config + + assert resolve(seed, tuning.TABLE_TACTIC, 8, MAX_BLOCKS) is seed + assert resolve(seed, None, 8, MAX_BLOCKS) is seed + # Malformed. + assert resolve(seed, (1, 2), 8, MAX_BLOCKS) is seed + assert resolve(seed, "nonsense", 8, MAX_BLOCKS) is seed + assert resolve(seed, (99, 1, 128), 8, MAX_BLOCKS) is seed + # Stale: tuned against a larger workspace than this one was built with. + assert resolve(seed, (0, 128, 128), 8, 32) is seed + # Stale: a variant this world size does not dispatch. + assert resolve(seed, (int(IpcVariant.FLAT_STAGED), 1, 128), 4, MAX_BLOCKS) is seed + # Usable. + assert resolve(seed, (int(IpcVariant.STAGED_RING), 2, 256), 8, MAX_BLOCKS) == ( + IpcLaunchConfig(2, 256, IpcVariant.STAGED_RING) + ) + + +def test_cache_key_extras_separate_every_workspace_dimension() -> None: + """The autotuner's own key is only the bucketed shapes. + + Without these, a TP4 and a TP8 entry at the same shape share one slot, a + configuration measured on one fabric is reused on the other, and a cache + written for one workspace size is applied to another. + """ + base = dict( + world_size=8, + profile=PROFILE_ROOTCPLX, + max_blocks=128, + max_numel=6144 * 128, + dtype=torch.bfloat16, + ) + reference = tuning.cache_key_extras(**base) + assert reference == tuning.cache_key_extras(**base), "must be deterministic" + assert isinstance(hash(reference), int), "extras must be hashable" + + for field, other in ( + ("world_size", 4), + ("profile", PROFILE_SWITCHPAIR), + ("max_blocks", 32), + ("max_numel", 6144 * 64), + ("dtype", torch.float16), + ): + assert tuning.cache_key_extras(**{**base, field: other}) != reference, field + + +def _synthesise_cache(monkeypatch, *entries) -> None: + """Install file-backed entries keyed the way ``save_configs`` writes them.""" + from flashinfer.autotuner import AutoTuner + + configs = { + str((tuning.PCIE_IPC_CUSTOM_OP, "PcieIpcAllReduceRunner", ((1, 4096),), e)): [ + "PcieIpcAllReduceRunner", + [1, 8, 256], + ] + for e in entries + } + monkeypatch.setattr(AutoTuner.get(), "_file_configs", configs, raising=False) + + +def test_a_cache_written_for_another_workspace_reads_as_uncovered(monkeypatch) -> None: + """A key mismatch misses every entry, which looks exactly like a hit.""" + tuned = dict( + world_size=4, profile=PROFILE_ROOTCPLX, max_blocks=128, max_numel=6144 * 128 + ) + _synthesise_cache( + monkeypatch, tuning.cache_key_extras(**tuned, dtype=torch.bfloat16) + ) + + assert tuning.cache_covers_workspace(**tuned) + for field, other in ( + ("max_numel", 6144 * 256), + ("world_size", 8), + ("max_blocks", 32), + ("profile", PROFILE_SWITCHPAIR), + ): + assert not tuning.cache_covers_workspace(**{**tuned, field: other}), field + + +def test_coverage_ignores_the_dtype(monkeypatch) -> None: + """One workspace serves both 2-byte dtypes, each with its own entries. + + Requiring a dtype match would report an uncovered workspace for a cache that + covers it perfectly well in the dtype the caller is not using right now. + """ + tuned = dict( + world_size=4, profile=PROFILE_ROOTCPLX, max_blocks=128, max_numel=6144 * 128 + ) + _synthesise_cache( + monkeypatch, tuning.cache_key_extras(**tuned, dtype=torch.float16) + ) + assert tuning.cache_covers_workspace(**tuned) + + +def test_an_empty_cache_covers_nothing(monkeypatch) -> None: + _synthesise_cache(monkeypatch) + assert not tuning.cache_covers_workspace( + world_size=4, profile=PROFILE_ROOTCPLX, max_blocks=128, max_numel=6144 * 128 + ) + + +def test_cache_key_extras_are_synthesis_invariant() -> None: + """The tuner keys on synthesized tensors and looks up with real ones. + + Anything derived from tensor *content* would make those two disagree and + every lookup would miss. + """ + real = torch.zeros(4, 6144, dtype=torch.bfloat16) + synthetic = tuning.small_int_initializer( + (4, 6144), torch.bfloat16, torch.device("cpu") + ) + common = dict( + world_size=8, profile=PROFILE_ROOTCPLX, max_blocks=128, max_numel=6144 * 128 + ) + assert tuning.cache_key_extras(dtype=real.dtype, **common) == ( + tuning.cache_key_extras(dtype=synthetic.dtype, **common) + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_synthesized_inputs_reduce_exactly(dtype: torch.dtype) -> None: + """Zero-tolerance verification is only valid if the group sum is exact. + + The kernels sum in a different order than NCCL, so an inexact sum would + show up as a mismatch and the gate would reject every candidate. + """ + max_sum = (tuning.INIT_MAX_VALUE - 1) * max(_WORLD_SIZES) + assert max_sum <= 256, "must stay inside bfloat16's exact-integer range" + + values = torch.arange(0, tuning.INIT_MAX_VALUE, dtype=torch.int32).to(dtype) + for world_size in _WORLD_SIZES: + summed = (values.float() * world_size).to(dtype) + assert torch.equal(summed.float(), values.float() * world_size) + + synthetic = tuning.small_int_initializer((8, 64), dtype, torch.device("cpu")) + assert synthetic.dtype is dtype + assert int(synthetic.float().min()) >= 0 + assert int(synthetic.float().max()) < tuning.INIT_MAX_VALUE + assert (synthetic.float() == synthetic.float().round()).all() + + +def test_verdict_reduces_with_max_over_wrong() -> None: + """The verdict must be a group decision, and MAX over "was wrong" is it. + + Corruption in this protocol is not uniform: a rank can be clean while its + peers are wrong. A rank-local verdict, or a SUM, lets a clean rank keep a + candidate its peers rejected -- and the ranks then profile different + candidate sets, which deadlocks the autotuner's timing reduction on the + first divergence. + """ + calls = [] + real_dist = tuning.dist + try: + tuning.dist = type( + "D", + (), + { + "all_reduce": staticmethod( + lambda tensor, op=None, group=None: calls.append(op) + ), + "ReduceOp": real_dist.ReduceOp, + }, + ) + tuning.reduce_verdict(torch.zeros(3, dtype=torch.int32), None) + finally: + tuning.dist = real_dist + + assert calls, "the verdict must be reduced, not decided locally" + assert all(op is real_dist.ReduceOp.MAX for op in calls), ( + f"a wrong-flag verdict must reduce with MAX (logical OR), got {calls}" + ) + + +def test_tuning_config_is_shared_and_its_mapper_is_stable() -> None: + """The lookup and the search must map a shape through the identical mapper. + + ``_find_nearest_profile`` is memoised on the spec, so a freshly built + mapper on every call would both miss the cache and grow it without bound. + """ + assert tuning.pcie_ipc_tuning_config() is tuning.pcie_ipc_tuning_config() + buckets = (1, 2, 4) + assert make_bucket_mapper(buckets, round_map=False) is ( + make_bucket_mapper(buckets, round_map=False) + ) + config = tuning.pcie_ipc_tuning_config() + (spec,) = config.dynamic_tensor_specs + assert spec.input_idx == (0,) and spec.dim_idx == (0,), ( + "only the batch dimension may bucket; hidden must stay exact in the key" + ) + assert config.constraint_specs == (), ( + "a constraint dim is stored as -1, which would erase hidden from the key" + ) + + +def test_batch_buckets_never_exceed_the_batch_they_stand_for() -> None: + """Floor semantics: a bucket is always a batch that was actually measured. + + Rounding up would apply a configuration measured at 64 to a batch of 33. + """ + mapper = make_bucket_mapper(tuning.TUNE_BATCHES, round_map=False) + for batch in range(1, 257): + bucket = mapper(batch) + assert bucket in tuning.TUNE_BATCHES + assert bucket <= batch or batch < min(tuning.TUNE_BATCHES) + + +def test_workspace_capacity_filters_buckets() -> None: + """A bucket larger than the workspace would raise inside the collective.""" + assert tuning.tuned_batches_for(6144, tuning.TUNE_BATCHES, 6144 * 8) == ( + 1, + 2, + 4, + 8, + ) + assert tuning.tuned_batches_for(6144, tuning.TUNE_BATCHES, 6144 * 128) == ( + tuning.TUNE_BATCHES + ) + + +def test_custom_op_name_is_stable() -> None: + """It is baked into every persisted cache key; renaming it orphans the file.""" + assert tuning.PCIE_IPC_CUSTOM_OP == "flashinfer::pcie_ipc_all_reduce" + assert tuning.PCIE_IPC_TUNE_VERSION == 1 + + +def test_pack_config_is_injective_over_the_candidate_space() -> None: + """The cross-rank agreement check compares packed configurations.""" + packed = {} + for world_size in _WORLD_SIZES: + for tactic in tuning.candidate_tactics(world_size): + config = tuning.tactic_to_config(tactic) + key = tuning.pack_config(config) + assert packed.setdefault(key, config) == config + + +def test_unrelated_autotune_context_preserves_workspace_cache(monkeypatch) -> None: + """Global tuning without a matching group must look up, not profile. + + A caller may wrap model warmup in an autotune context that replaces the + singleton tuner's file cache without installing the process group required + to profile a collective safely. The workspace must restore its explicit + cache instead of resolving and retaining the seed tactic. + """ + + class Runner: + def can_profile(self, device) -> bool: + return False + + class Tuner: + is_tuning_mode = True + + def __init__(self) -> None: + self.loaded = [] + self.lookups = 0 + + def load_configs(self, path: str) -> None: + self.loaded.append(path) + + def search_cache(self, *args, **kwargs): + self.lookups += 1 + return True, 0, (int(IpcVariant.STAGED), 16, 256), None + + def choose_one(self, *args, **kwargs): + raise AssertionError( + "collective profiling must not run without a tune group" + ) + + workspace = PcieIpcAllReduceWorkspace.__new__(PcieIpcAllReduceWorkspace) + workspace._runner = Runner() + workspace._tune_batches = tuning.TUNE_BATCHES + workspace._tune_cache = "/tmp/pcie-ipc-tuned.json" + workspace._tune_cache_exists = True + workspace._tuned_configs_loaded = True + workspace._warned_untuned = False + workspace.world_size = 2 + workspace.max_blocks = MAX_BLOCKS + workspace.group = object() + workspace.device = torch.device("cpu") + + monkeypatch.setattr( + "flashinfer.comm.pcie_ipc_ar.dist.all_reduce", lambda *a, **k: None + ) + tuner = Tuner() + seed = IpcLaunchConfig(8, 64, IpcVariant.UNSTAGED) + result = workspace._resolve_tuned(torch.empty(4, 4096), seed, tuner) + + assert tuner.loaded == [workspace._tune_cache] + assert tuner.lookups == 1 + assert result == IpcLaunchConfig(16, 256, IpcVariant.STAGED) diff --git a/tests/trace/template_registry.py b/tests/trace/template_registry.py index 407d1f5efb0..b42bc1d22fd 100644 --- a/tests/trace/template_registry.py +++ b/tests/trace/template_registry.py @@ -41,6 +41,7 @@ "flashinfer.cascade", "flashinfer.comm.allreduce", "flashinfer.comm.dcp_alltoall", + "flashinfer.comm.pcie_ipc_ar", "flashinfer.concat_ops", "flashinfer.cudnn.decode", "flashinfer.cudnn.prefill",