Repository navigation
Add opt-in clustered phased PCIe all-reduce - #70
Conversation
Prepare fixed 64-CTA, 512-thread joint launches with one active eight-CTA cluster for phased BF16 transfers. Preserve ordinary profile identities and IPC layouts while validating cluster eligibility and residency. Validate profile admission and owner intervals on CPU, and exact CUDA graph replay, startup selection, aliasing, allocation stability, and profile transitions on TP4 SM120. AI-assisted implementation and review.
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Cluster profile contract and plan admission flashinfer/experimental/b12x/comm/pcie/_joint_profile.py, pcie_allreduce.py, _tuning.py, tests/experimental/b12x/preparation/test_joint_clusters.py, README.md |
Cluster size 8 requires the documented launch geometry and hardware constraints. The profile registry adds phased modes, and the public plan and query validation enforce where cluster metadata is accepted. |
Cluster kernel and prepared launch path flashinfer/experimental/b12x/comm/pcie/_joint_cluster_cute.py, _joint_preparation.py, tests/experimental/b12x/preparation/test_joint_cluster_intervals.py, test_joint_clusters.py |
A cluster launcher implements phased and ordinary dispatch. Preparation selects it for clustered plans and validates geometry and residency. Tests cover owner intervals, compile jobs, and materialization. |
Startup profile selection and tuning flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py, _joint_startup.py, tests/experimental/b12x/preparation/test_joint_clusters.py, README.md |
Startup profile validation, cache keys, and collective tuning account for cluster size and its available profiles. |
Graph replay and serving behavior tests/experimental/b12x/comm/test_joint_cluster_gpu.py, flashinfer/experimental/b12x/comm/pcie/README.md |
Opt-in GPU tests cover graph launch geometry, profile transitions, outputs, and startup selection. The documentation describes cluster launch behavior and the shared compiled launcher. |
Estimated code review effort: 4 (Complex) | ~60 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant Caller
participant PCIeAllReduce
participant JointPreparation
participant JointClusterLauncher
Caller->>PCIeAllReduce: plan with joint_cluster_size=8
PCIeAllReduce->>JointPreparation: validate geometry and prepare plan
JointPreparation->>JointClusterLauncher: compile launcher and validate residency
Caller->>JointClusterLauncher: replay graph with selected profile
Possibly related PRs
- local-inference-lab/flashinfer#62: It defines overlapping phased mode identifiers, so the profile registries and mode semantics need reconciliation.
Suggested reviewers: lukealonso, aleozlx
Merge Risk: 
560fe
-
No actionable merge-blocking issue remains in the reviewed change; proceed with normal validation.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 19.35% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title "Add opt-in clustered phased PCIe all-reduce" matches the main change in the changeset. The diff adds cluster-scoped phased profiles (full_phased, payload_rsag_phased) to the experimental b1… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Description check The pull request description follows the required template and provides detailed scope, related issues, testing evidence, experimental-track information, declared test targets, reviewer notes, and kno… Full details: Title check
-
Explanation
The title "Add opt-in clustered phased PCIe all-reduce" matches the main change in the changeset. The diff adds cluster-scoped phased profiles (full_phased, payload_rsag_phased) to the experimental b12x PCIe joint kernels, gated by joint_cluster_size=8 and explicit opt-in. The title is concise, specific, and free of file lists or vague terms.
Full details: Description check
-
Explanation
The pull request description follows the required template and provides detailed scope, related issues, testing evidence, experimental-track information, declared test targets, reviewer notes, and known CI limitations. The manual repository-wide pre-commit check remains unchecked, and remote GPU CI did not run, but these omissions are clearly disclosed and do not make the description incomplete.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 19.35% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 31 functions across 10 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate docstrings
-
- Commit to this branch
- Create a new PR
Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
|
@flashinfer-bot run |
📌 Description
Add explicitly selected, cluster-scoped phased BF16 PCIe all-reduce profiles to the experimental b12x joint kernels.
PCIeAllReduce.plan(..., joint_only=True, joint_cluster_size=8)prepares a fixed 64-CTA grid with 512 threads per CTA and eight CTAs per cluster. The phased profiles use one active cluster to publish peer destinations in order, with a system fence from every writer followed by a cluster barrier between remote destinations. This constrains peer traffic without a cooperative whole-grid launch.The latency family admits
full_phasedandpayload_rsag_phased; the bandwidth family admitspayload_rsag_phased. Preparation rejects unsupported architecture, geometry, capacity, profile, and compiled residency. Device profile selection, captured startup tuning, workspace ownership, BF16 arithmetic, IPC layout, and native launcher arguments retain their established contracts. Ordinary profiles remain available inside the cluster launcher; they retain cluster launch geometry, so they carry no zero-overhead guarantee.Compatibility: omitted
joint_cluster_sizeand explicit1retain the ordinary query fields, profile encodings, launch policy, and compile-job identities. The native preparation query schema advances from 41 to 42 so preparation caches distinguish the expanded query contract. Experimental APIs have no compatibility guarantee. No AOT registration or automatic stable dispatch is added.Implemented: explicit clustered plans, phased profiles, residency validation, and startup profile selection. Qualified: TP4 on NVIDIA RTX PRO 6000 Blackwell Workstation Edition with CUDA 13.2 for both joint families and aliased/disjoint outputs. Other admitted rank counts have CPU partition coverage but no additional GPU qualification in this PR validation run. Unsupported: clustered geometry outside SM120/SM121 with at least 64 SMs, grids other than 64 CTAs × 512 threads, and phased active counts other than eight.
The serving comparison uses an ordinary-ladder package and a package prepared with
gb202_pcie_tp4_clustered_full.json. Serving measurements used DeepSeek V4.1 Flash 1017, TP4/C32/K7, maximum 128 verification rows, prefill 4096, native expert scales, BF16 MLA, RAM Engram tables and a 64 GiB prefix cache. The two machines, frank1 and frank2, have AMD EPYC 9575F CPUs and RTX PRO 6000 Blackwell 96 GB GPUs. Memory offsets were +6000 MHz; graphics clocks remained automatic. Both packages used the balanced scheduler, temperature/top-p 1 and, in this table, alpha 0.6.Each host-bridge value is the mean of five 30-second LLM Bench cells on GPUs 0,4,8,12. Each switch value combines three cells on 0–3 and three on 4–7. Switch measurements ran two active services per host; host-bridge measurements ran one. These are within-placement comparisons, not an isolated measurement of topology effects. Cells used context 0/max 8192 output tokens and shared adaptive state within each native process; no IID or fixed-work GPU-latency claim is made.
All 352 accepted samples, including alpha 0.8 retain emitted-token throughput, completed steps/s, maximum-rank GPU time per step, verified draft depth and width-one fraction. Host-bridge C8/C16 improved across both machines; switch/C1 cases include lower means, so the clustered catalog remains explicit opt-in. All ordinary profiles remain available, but their compiled cluster geometry is still active.
Two synchronized groups of four cells were excluded and rerun with unchanged clients/guards: a C1 request-rollover gap violated full concurrency, and one C32 failure counter increased after the timed client finished. The client reported zero errors; an existing Flare unit test confirms a closed output channel can increment that counter, but the incident-specific error was not logged.
The serving packages used Flare 6fa01a49 and FlashInfer 361e45f plus these clustered changes. Full-model performance on the Flare PR base 881eccdd has not been rerun. The signed candidate runtime contains the measured preparation implementation; promotion preserved its execution descriptors and it passed 192 operational requests in the Flare 6fa Docker image.
🔍 Related Issues
Experimental tracking: voipmonitor/flare#7. Owner: @voipmonitor.
Companion preparation and native serving measurements: NVIDIA Flare #174.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.All configured hooks applicable to the eleven changed files passed; the installed pre-commit hook also runs at commit time. Repository-wide formatting was not run.
🧪 Tests
unittest, etc.).Focused CPU validation passed 140 tests across
test_joint_clusters.py,test_joint_cluster_intervals.py,test_joint_profiles.py,test_joint_covering_grids.py,test_collective_family_buckets.py, andtest_collective_tuning.py, all undertests/experimental/b12x/preparation/.The four cases in
tests/experimental/b12x/comm/test_joint_cluster_gpu.pypassed on four RTX PRO 6000 Blackwell GPUs using CUDA 13.2. Conditions include both families, aliased/disjoint outputs, twelve payload sizes from 16 to 2,293,760 bytes, ordinary/phased profile transitions, fixed CUDA graph launch attributes, epoch rollover, poisoned reusable storage, exact numerical references including special values, allocation stability during repeated replay, malformed active-count selectors, and captured startup selection with rank agreement. Every tested ordinary and phased result matched its numerical contract exactly.Reproduce the GPU suite with a four-GPU SM120 allocation with at least 64 SMs per GPU and the experimental b12x compatibility import path:
CUDA_VISIBLE_DEVICES=0,1,2,3 B12X_RUN_JOINT_CLUSTER_TEST=1 \ PYTHONPATH="$PWD:$PWD/flashinfer/experimental/b12x/_compat" \ python -m pytest tests/experimental/b12x/comm/test_joint_cluster_gpu.py -q🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: Track qualification of clustered phased PCIe collectives voipmonitor/flare#7tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an@flashinfer_experimental_apior naming a backend explicitly is itself the opt-in and needs no environment variable.)Reviewer Notes
The cluster implementation intentionally preserves the measured kernel geometry and numerical protocol. Review should focus on cluster-uniform barrier participation, per-thread system fences, complete owner intervals for uneven/tiny payloads, and immutable launch metadata through preparation and startup tuning. Ordinary profiles on an explicitly clustered plan use the same cluster launch attributes even though they do not execute cluster barriers. CUDA graph inspection verifies no cooperative launch attribute and the intended cluster dimension.
The runtime parameter is explicit opt-in within an already experimental backend. No core files change. The README includes a plan declaration and a startup profile menu. Whole-grid cooperative synchronization is outside this implementation.
Owner @voipmonitor proposes graduation review by 2026-11-30, targeting the first FlashInfer release after maintainers approve the geometry and API contract. Graduation requires finalized public configuration, supported-hardware CI coverage, broader rank-count GPU qualification, and reproducible topology-specific performance evidence. Until those conditions are met, automatic stable dispatch remains excluded.
CI observations for head
560fef16648e: the documentation checks and docs build pass. The repository-wide formatter reports ten files outside this PR's changed-file set; changed-file hooks passed locally. The labeler has no authentication token, and the CI bot rejected the test request becausevoipmonitoris not authorized. Remote GPU CI has therefore not run. These checks are not reported as passing, and this PR does not change repository authorization or unrelated formatting.