You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add serialized PCIe profiles and interleaved startup sampling - #62
PCIe reduction tuning needs to compare different peer-write schedules without changing reduction arithmetic or requiring every CTA to be resident. It also needs samples spread across the prepared collective groups, rather than consecutive measurements of one group. Destination serialization and interleaved startup sampling address those two constraints within the existing prepared joint BF16 kernels.
Relative to FlashInfer 0b6178857, the intended change adds two profiles: serial_push batches writes by destination in one active CTA; queued_push distributes a destination's tiles among CTAs and advances after completed tiles, so progress does not depend on unlaunched CTAs. Neither profile is admitted at TP2. Exact input/output aliases use coherent loads; disjoint bulk inputs retain their read-cache behavior. The owner-cyclic FP32 reduction and single BF16 rounding remain unchanged. Static finalist capacity increases from 35 to 55; kernel compile key 20 distinguishes the implementation.
Startup tuning changes from consecutive per-group sampling to 21 interleaved passes. Each candidate is scored by the sum of complete-chain maximum-rank durations. Each replay resets status, validates arithmetic, and rejects invalid or overflowing scores. Prepared profile encoding and the native serving contract remain unchanged within this intended collective subset.
The expected benefit is a broader choice of safe write schedules and samples separated by other groups. A performance improvement is not qualified: the switched-fabric regression below blocks that claim. This PR remains draft.
Scope limitation: the intended collective changes occupy 19 files relative to 0b6178857, but GitHub's comparison of head 729464839 against base e16f9495 displays 848 files after upstream history was rewritten. That broader comparison includes compiler, attention, GEMM/MoE, and core CUDA changes. The historical qualification below covers the collective subset, not the entire displayed PR. The broad comparison must be resolved and qualified before merge.
🔍 Related Issues
Flare #128 captures the prepared producers into one startup graph, binds their selectors into serving graphs, and checks aggregate status during native loading. It depends on this API; FlashInfer integration must precede the Flare dependency update.
🚀 Pull Request Checklist
✅ Pre-commit Checks
I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.
Changed-file hooks passed for the intended collective changes. No full-comparison hook result is established.
🧪 Tests
Tests have been added or updated as needed.
All tests are passing (unittest, etc.).
Qualified collective correctness: 204 focused CPU tests, 11 GPU collective cases for compile key 20, and GPU score-publication checks passed. Coverage includes aliased/disjoint output, FP32 addition order, signed zeros, tails through 4096 rows, profile changes, epoch reuse, allocation stability, and observable protocol progress. GPU qualification is TP4 on SM120; other admitted rank counts are unqualified. Four restricted-residency cases passed for the unchanged queue protocol before the bulk alias-load correction; those four results do not constitute a rerun of the corrected implementation.
An additional dispatch/registry run passed 74 tests and failed test_registry_matches_disk because quantization.mla_norm_mxfp8 was absent from _OPS. The same failure reproduced on clean 0b6178857.
Performance blocker: on switched RTX PRO 6000 SM120 TP4, payload reduce-scatter/all-gather (payload_rsag) at BF16 [31,5120], 512 threads and 64 active CTAs measured 12.244 → 14.552 µs p50 (+18.9%). The affected range is approximately 17–32 rows. Sustained eight-call probes used at least two seconds of warmup and three seconds of timing, with repeated upstream controls. A 3090 MHz graphics clock was requested; observed SM clocks were 2842–2910 MHz. Same-cubin 0/16-byte shared-reservation measurements were flat, and dispatcher reordering did not remove the regression. Its cause remains unresolved.
Whole-model qualification of this change is pending. For context, upstream FlashInfer 0b6178857 with Flare eb71702b, DeepSeek V4.1 Flash package gb202-tp4-202610290017, measured 190.35/345.74/391.53 tok/s at C1/C8/C16 across host bridges and 319.29/869.97/1162.69 tok/s on a switched group. The native-static host-bridge reference measured 247.90/463.16/511.29 tok/s. These are upstream controls, not gains from this PR. They used RTX PRO 6000 SM120 TP4, prepared concurrency 16, fixed width 8, 64G prefix cache, two 30-second cells per concurrency, unmodified boost clocks, Engram weights on disk, scales in RAM, and a 4 GiB row cache. RAM-delivery comparisons are unqualified. Flare #128 records the additional 90b5c6d3 control and its differing kernel catalog.
Remaining gates include regression attribution and correction, measured ladder additions, full-model C1/C8/C16 checks, three native starts, and two independent exports. Interleaving one group alone cannot provide separation by other groups.
The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.
The core-scope, test-placement, and registration checks are unconfirmed for the 848-file comparison. The hardware results and runnable collective benchmark apply to the intended subset. The targeted CI fence does not cover the broader diff.
The serialized profiles retain the prepared physical grid and update only the admitted mode/active-CTA selector. Offline runtime ladders bypass the preparation-time geometry race. Review of this intended mechanism does not establish correctness or compatibility of the unrelated files in the displayed comparison. Merge remains blocked on that scope mismatch and the performance gates above.
Port runtime, tests and benchmarks from b12x 64836b0e4 through ffabbadee3.
Add prepared paged decode and KV writes, BF16 and block-FP8 GEMV paths,
TP3 PCIe collectives, compressed W4A16 scale layouts and residency fixes,
FP32 router-weight accumulation, and preparation cache updates.
Validate paged-attention buffer layouts and devices, preserve scratch
ownership, and enforce block-FP8 capacity, operand and alpha contracts.
Use the shared preparation sampler and in-tree references in benchmarks.
Preserve FlashInfer packaging and the exclusion of vLLM adapters.
Validation on RTX PRO 6000 Blackwell (SM120): paged decode 55 passed;
block-FP8 GEMV 40 passed; CSF residency 4 passed across capacities 1-256
and 8192; mHC PDL 8 passed; TP3 prepared oneshot and DMA passed, including
oneshot with PDL. Attention/GEMV and MoE/preparation regression runs,
BF16/FP8 attention benchmarks and six block-FP8 benchmark cases completed.
All 6737 b12x tests collect and applicable pre-commit hooks pass.
The broader W4A16 run has the same 70 failures and 16 skips as the base
revision. The legacy oneshot torture test also fails on both revisions
because it omits the required preparation plan.
Compile the direct reader before the batch reader, whose bounce implementation depends on IO_ALIGNMENT and validate_direct_range. Protect the implementation includes from clang-format sorting.
Validation: native helper compilation, five direct-I/O tests, and pre-commit checks passed. AI-assisted.
Reuse packed NVFP4 expert weights for per-call NVFP4 activation prefill,
retaining calibrated scales only for eligible A16 plans. Preserve the
inclusive A16 cutoff, CSF scale preparation, and the canonical clamped
SiLU BF16 rounding contract. Keep hybrid disabled by default.
Snapshot threshold, activation-plane count, warp count, and router precision
in tuning query schema 23. Retain compiled launchers through the device-owned
compiler cache and caller-provided scratch for graph replay. Port the hybrid
behavior from b12x PRs 472, 473, and 476.
Validate 122 focused preparation, quantization-contract, selection, CSF,
and graph tests on SM120; the two-device ownership test passed separately.
Validate E288/H4096/I1024/topk8/1536 CSF with clamp10 against the float64
quantization oracle and changed-input graph replay. Model quality remains
a deployment qualification requirement.
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Normalize A4 terms and warps to neutral defaults when the prefill threshold disables the feature. Validate explicit controls only for enabled hybrid plans, so stale A4-only environment values cannot block ordinary MoE preparation.
Validation: five preparation-control tests passed with CUDA hidden; applicable pre-commit hooks passed. Enabled kernel arithmetic is unchanged.
Reconstruct packed scale bytes in shared-memory stages while retaining one resident weight payload and the dense scale path. Prepare hybrid capability with a boolean control and select precision explicitly at bind time. Share prepared-variant and scale-consumer selection with the NVFP4 prefetch query; use a larger prepared variant when an explicit A4 launch needs its scratch capacity.
Validation: 165 tests passed in the combined changed suites, with the two-device ownership case separately passing on GPUs 8 and 9. Pre-commit passed. Actual GLM TP2 layer inputs at M128, M512, M1536, and M3072 matched baseline A4 outputs bitwise and passed changed-input graph replay and allocation checks.
Prepare stage and dense A4 scale consumers without duplicating weight payloads or scratch. Keep stage decoding by default and for live row counts at or below the immutable CSF stage limit. Select dense consumption only when the caller guarantees ordered expanded scales. Share readiness and size selection between binding and the NVFP4 scale-consumer query; precision remains explicit per call.
Validation: 115 host precision tests and eight GPU correctness cases passed. GPU cases cover reader selection at 1536/1537 rows, dense/stage bitwise parity, changed-input CUDA graph replay, ordered expansion, and allocation stability. Pre-commit passed. Serving performance qualification remains separate.
Mask the final FC1 half tile and stage gate/up CSF half-slabs using their original exception records. Keep the aligned branch and numerical contract unchanged, without padding or duplicating weight payloads. Bump the compiler identity and tuning query schema.
Validate dense and CSF readers with one- and two-plane quantization, changed-input graph replay, payload aliasing, and byte-exact half-slab reconstruction. The float64 oracle also covers synthetic H2560/I320/top-k10 rank geometry.
Remove import-time PTX assembler discovery and environment mutation so Triton planning and child-process compilation retain the same compiler identity. Document explicit Blackwell assembler selection before process startup.
Validation: CUDA_VISIBLE_DEVICES=1 /home/luke/projects/vllm-jj-rebase/.venv/bin/python -m pytest tests/test_env.py -q (8 passed); pre-commit run --files docs/installation.rst flashinfer/triton/__init__.py tests/test_env.py (passed).
Assisted-by: OpenAI Codex
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Expose b12x plans for cache writes with separate key/value dimensions and raw-input depthwise convolution with caller-owned history and scratch. Share convolution kernels with PLE and retain graph-safe dynamic metadata.
Carry an optional fixed attention partition count through QSA planning and execution. Round raw-ring capacity to a page divisor that covers the rollback horizon, including seven drafts on 64-token pages. Add capture and rollback regressions.
Propagate deterministic output requirements through workspace planning and tuning legality. Invoke compiled Triton router launchers using their full positional signature.
Include recurrence policy and intermediate precision in prepared query identity. Keep BF16 chunk execution as the default; provide an explicit FP32 device recurrence for numerical diagnostics and a sequential single-key-split policy for callers requiring fixed reductions.
Retain prepared attention with selectable intermediate arithmetic for identical-input comparisons. Preserve the production default and isolate the eager-BF16 diagnostic implementation in preparation identity.
Split normalization sums and probability writes across CTAs while preserving each logical reduction lane and the 1024-thread final reduction. Reuse histogram storage after its final radix pass.
curand_uniform can return exactly one. Accept positive target-to-draft ratios of at least one in both chain acceptance loops while retaining RNG draws and the strict lower-ratio comparison.
Regression cases cover equal distributions, point masses, independent acceptance counts, lower ratios, scalar/device RNG, and graph replay. Targeted GPU qualification: 72 passed.
…e norm
Allow exact in-place residual updates while rejecting unsafe output aliases. Specialize vectorized affine loads by gain dtype and centering, preserving BF16 residual rounding.
Validation: 36 prepared norm and compilation tests plus the direct CuTe live-count reuse test passed on SM120; changed-input graph replay covers capacities through 4096. Ruff checks pass.
Carry the input row stride through convolution and speculative-history updates so packed QKV/Z outputs require no intermediate copy. Preserve contiguous columns, disjoint mutable buffers, and exact declared strides.
Validation: 26 short-convolution tests passed on SM120, including changed-input graph replay in decode, prefill, and mixed modes with 6144 columns of padding. Ruff checks pass.
Make the prepared BF16 and FP32 paths opaque to torch.compile while declaring scratch, output, and recurrent-state writes. Rebind compiler-supplied tensors to the selected plan without changing kernels or precision.
Validation: four fullgraph compilation and CUDA replay cases pass across both precisions and 17/65-token inputs, including changed operands, empty requests, checkpoints, exact state/output equality, and allocation-stable replay.
Declare selected logits, IDs, and route weights as custom-op mutations while retaining prepared scoring, bias, and modality contracts. Validate caller options against the selected invocation before dispatch.
Validation: all 13 routing tests pass, including fullgraph compilation and changed-input CUDA replay for both score functions, strided outputs, modality controls, and allocation stability.
Register precise output and barrier mutations for expert-presence masks, NVFP4-CSF expansion, and deterministic route reduction. Retain supplied scale-decoder programs and resolve the reduction stream at invocation.
Validation: 131 GPU tests passed; one two-device validation test skipped on the single assigned device. Includes compiled captures, changed routes, poisoned outputs, and allocation-stable replay.
Register output-mutating operators around softmax, top-k masking, top-p renormalization, and probability sampling. Preserve kernel selection, RNG tensors, and public return values.
Validation: 196 sampling pipeline and RNG tests passed on SM120, including full-graph compilation, CUDA replay, changed parameters, and signed scalar and uint64 tensor seeds.
Expose the prepared writer through a mutating Torch operator, preserving runtime tensor binding and validation. BF16 and FP8 strided cache writes remain visible to ATen tracing.
Validation: 13 SM120 writer tests passed, including compiled graph replay with changed slot mappings and strided packed K/V buffers.
Declare the output tokens and acceptance counters as mutable outputs of the existing native verifier. Preserve scalar and device RNG inputs, public signatures, and kernel arithmetic. Sampling compilation, capture/replay, and RNG coverage: 208 CUDA tests pass.
Register the oneshot and fused launcher memo tables with program_cache so deferred compilation closures are discarded when preparation resolves their kernel programs.
Validation: PCIe preparation tests (28 passed, 9 CUDA cases skipped); Flare compiled TP2 and TP4 native graph replay passes for a 4096-element BF16 reduction.
No actionable comments were generated in the recent review.
Walkthrough
The pull request adds phased variants to selected PCIe joint collective modes. These modes publish data to peer destinations in rank-staggered order, with system fences and phase synchronization. Profile-driven mode and thread lists now feed tuning, benchmarks, and tests.
Changes: Phased PCIe joint collectives
Layer / File(s)
Summary
Mode definitions and phase synchronization flashinfer/experimental/b12x/comm/pcie/_joint_profile.py, flashinfer/experimental/b12x/comm/pcie/_cute_intrinsics.py, flashinfer/experimental/b12x/comm/pcie/_joint_cute.py
The profile defines phased mode IDs, families, thread counts, and signal storage. New intrinsics and phase state support system fences and CTA arrival and release.
Phased publication and bulk paths flashinfer/experimental/b12x/comm/pcie/_joint_cute.py
_ready and _bulk use rank-staggered destinations and phase waits for selected modes. Kernel dispatch passes phase state and selects the phased paths.
Tuning and validation integration flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py, flashinfer/experimental/b12x/comm/pcie/_joint_startup.py, tests/experimental/b12x/comm/*, benchmarks/experimental/b12x/benchmark_pcie_allreduce_transports.py, flashinfer/experimental/b12x/comm/pcie/README.md
Tuning and finalist sizing use profile-defined modes and thread counts. The benchmark and tests use profile-defined modes and strip the phased suffix for reference behavior. The README lists and describes the phased candidates.
sequenceDiagram
participant JointLaunch as _JointLaunch
participant Ready as _ready
participant Phase as _phase
participant Intrinsics as local_phase_arrive_and_wait
participant PeerSlabs as peer slabs
JointLaunch->>Ready: publish data to destinations
Ready->>Phase: fence and synchronize between destinations
Phase->>Intrinsics: wait for CTA arrivals at target
Intrinsics-->>Phase: release when all CTAs arrive
Ready->>PeerSlabs: write destination data
Loading
Suggested reviewers:lukealonso, aleozlx
Merge Risk: Minimal · up to ebab2
No actionable issue remains from this review; the change is ready to merge after normal checks.
Pre-merge checks: 2 · 3
Failed checks: 3 warnings
Check name
Status
Explanation
Resolution
Docstring Coverage
Docstring coverage is 15.38% 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.
Title check
The title "Add serialized PCIe profiles and interleaved startup sampling" does not match the changeset. The diff adds phased joint write modes (full_phased, payload_rsag_phased, bulk_push_phased), pha…
Rename the PR to describe the main change, for example "Add phased joint write modes for host-bridge PCIe fabrics". This matches the PR objectives (PR #62) and the code changes.
Description check
The description is substantive and tracks the template well. It states the purpose (destination serialization and interleaved startup sampling for PCIe reduction tuning), lists the two new profiles, a…
1) Resolve the scope mismatch: the PR shows 848 changed files against base e16f949 while the intended collective change touches 19 files. Rebase or re-create the PR so the diff contains only the intended changes, then update the descriptio…
Passed checks: 2 passed
Check name
Status
Explanation
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.
Full details: Docstring Coverage
Explanation
Docstring coverage is 15.38% 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 13 functions across 8 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Full details: Title check
Explanation
The title "Add serialized PCIe profiles and interleaved startup sampling" does not match the changeset. The diff adds phased joint write modes (full_phased, payload_rsag_phased, bulk_push_phased), phase barriers, a system fence, and a startup finalist capacity change. It does not describe serialized profiles or interleaved startup sampling in a way the changed files support, and the PR objectives center on phased write modes for host-bridge PCIe fabrics. The title is therefore misleading about the primary change.
Resolution
Rename the PR to describe the main change, for example "Add phased joint write modes for host-bridge PCIe fabrics". This matches the PR objectives (PR #62) and the code changes.
Full details: Description check
Explanation
The description is substantive and tracks the template well. It states the purpose (destination serialization and interleaved startup sampling for PCIe reduction tuning), lists the two new profiles, and gives extensive test and performance evidence. It fills in Related Issues (Flare #128), the pre-commit section, Tests, and the Experimental Track section with tracking issue #62 and a declared test-scope fence. The author clearly marks the PR as draft and blocked. Remaining gaps: the pre-commit hooks box is unchecked, the 'All tests are passing' box is unchecked, and the Experimental Track sub-checklist (owner, graduation plan, thin entry point, runnable example) is unchecked. The description also says the GitHub comparison covers 848 files, far beyond the intended 19-file subset, which the author flags as unresolved. The 'Reason' line for full-scope CI is left as the template placeholder, and the CI target fence lists paths that the description does not confirm exist. These gaps mean the description is not fully complete, but it is not largely incomplete.
Resolution
Resolve the scope mismatch: the PR shows 848 changed files against base e16f949 while the intended collective change touches 19 files. Rebase or re-create the PR so the diff contains only the intended changes, then update the description to match. 2) Run pre-commit run --all-files on the corrected diff and check the box only after fixing reported issues. 3) Resolve the 'All tests are passing' item or state which tests fail and why, including the test_registry_matches_disk failure that also reproduces on clean 0b61788. 4) For the Experimental Track, complete the tracking issue #62 with an owner, the reason for the experimental path, and a graduation plan with a target release, or leave the boxes clearly unchecked with reasons. 5) Confirm that the listed experimental-tests targets exist under tests/experimental/ and that the Core changes stay within a thin entry point. 6) Replace the 'Reason:' placeholder for full-scope CI if full-scope CI is needed, or remove that line. 7) Resolve the unresolved performance regression (payload_rsag +18.9% at 17-32 rows on switched fabric) or document the decision to merge with it, since the description says merge is blocked.
Finishing Touches
Generate docstrings
Commit to this branch
Create a new PR
Fix failing CI checks
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.
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.
Use two KV MMAs per warp for 64-dimensional values, whose fragment layout rejects odd counts. Permit the resulting 128-row exact-plane BF16 stage: descriptor geometry, shared allocation, TMA byte counts and MMA iteration follow the selected tile. Retain the 64-row shape for other general decode dimensions.
Validation on NVIDIA SM120: 15 paged trait and public-preparation CUDA graph tests pass, covering BF16 H64/H128/H192, 128-token page boundaries, changing lengths and allocation stability. Flare compiled full/sliding H64 decode and complete continuation fork/restore tests also pass. No performance claim.
Expose caller-owned route outputs from the dynamic NVFP4 epilogue while retaining parallel accumulation over all intermediate slices. Prepare physical-order paired FC1 gains as owned kernel-order tensors and preserve calibrated activation scales.
Include output layout and gain geometry in preparation and compilation identities. Exclude fused token-only backends from route-output tuning while retaining legal dynamic tile candidates.
Validate route and default token outputs, decode and 4096-token prefill, graph replay with changed routes and poisoned scratch, scale-swap rejection, and every intermediate slice against an independent arithmetic interval reference. GPU tests: 20 passed; W4A8-MX regression tests: 10 passed. Checkpoint-dependent W4A8-NVFP4 tests skipped because the fixture checkpoint is unavailable.
Stop a paired race early only when every challenger sample exceeds twice the slowest champion sample. Keep adaptive balanced comparisons for close contenders and invalidate affected selection caches.
…lies
Use the largest declared plain all-reduce payload in each transfer bucket to measure every family. Retain the winning family and the fastest other family, then tune their eligible launch configurations for subsequent exact shapes. Separate calibration by rank group, dtype, aliasing, transport availability, and storage geometry. Include calibration choices in cached selection identity and retain unrestricted audit races.
Prepare head-dimension-64 paged decode and merge for BF16 queries with BF16
or E4M3 caches. Retain unweighted FC2 partials in FP32 scratch and reduce
them in fixed order before BF16 conversion. Use 64-bit flattened reducer
indices for output domains above two billion elements.
Validation on SM120: 32 captured head-64 attention cases passed with cosine,
numerical and allocation checks; 45 routed-expert and split-execution tests
passed, including eager and captured reduction beyond the int32 boundary.
Measure joint BF16 PCIe profiles across 209 payload sizes with a 600-second
charged budget, payload-scaled sampling, bracketing anchor coverage, and
balanced confirmation. Retain raw timings and topology in resumable JSONL
reports and support independent recommendation replay.
Render each report's kernel choices and exact bucket configurations to PNGs,
either after collection or through a CPU-only plot command. Matplotlib is
optional. Serving dispatch and preparation defaults remain unchanged.
Validation: 169 CPU tests passed; choice and bucket plots rendered from both
saved four-GPU topology reports. Prior GPU collection covered all 209 sizes
on both topologies. Independent replay of the refined separate-switch
recommendations remains unrun; no universal 5% regret bound is established.
Check collective eligibility and finalist policy before accessing tuning configuration. Keep finalist requests outside family pruning and document their full catalog coverage.
Validation: 171 CPU tests passed across collective tuning, family buckets, contextual APIs, offline ladder search and plotting, and DMA capture lifetime. No GPU benchmarks were run.
Keep one native candidate per logical PCIe plan and key selector memory by
that logical owner. Revisited candidates preserve compile jobs and separate
serving selectors without retaining every losing declaration.
Implemented for PCIe collective declarations. Qualified with 99 focused CPU
tests and a DeepSeek V4.1 TP4 runtime export paired with Ignition releasing
tuning owners before terminal collection. On frank2's four RTX PRO 6000
SM120 GPUs, release completed in 15.8 s under the 60 s supervisor deadline.
Prepared kernel execution and public operation contracts remain compatible.
Authored-by: Codex GPT-6
Add an autotuned joint profile that batches direct peer writes by destination
within one CTA. Batched coherent loads hide memory latency while preserving
owner-cyclic FP32 addition, one BF16 rounding, and existing bulk storage.
Admission requires at least three ranks and exactly one active CTA; thread
count remains autotuned. Mode identity has one Python enum and compile key 14.
Implemented; performance remains research-only. TP4 GPU checks cover aliases,
disjoint outputs, adversarial FP32 order, tails through 4096 rows, scratch poisoning,
epoch wrap, allocation stability, and observable protocol counters. Programs
have no local stack spills, but register usage increases. Full-model performance,
foreign-fabric startup stability, and switched regressions remain unqualified.
Authored-by: Codex GPT-6
voipmonitor
changed the title
Add phased joint write modes for host-bridge PCIe fabrics
Bound runtime collective metadata and batch peer writes by destination
Oct 28, 2026
Host-bridge all-to-all collapse: root cause is PCIe write ordering (WRX90 / Threadripper PRO)
We hit the host-bridge SM all-to-all collapse this PR is chasing ("research-only collective performance: the goal is not met") on two identical workstations. A short A/B points at strong ordering of GPU-issued peer TLPs. With the NVIDIA driver's relaxed-ordering mode on, the SM fabric all-to-all recovers 14–22x. b12x and p2pmark then stall, which suggests where a fence belongs.
Hosts: ASUS Pro WS WRX90E-SAGE SE (BIOS 1317), TR PRO 9965WX, 4x RTX PRO 6000 Blackwell Max-Q. Each GPU is Gen5 x16 on its own root port, so every pair crosses host bridges. Kernel 6.8.0-147, driver 615.71.09 with RegistryDwords="ForceP2P=0x11;GrdmaPciTopoCheckOverride=1;EnableResizableBar=1".
lil-bench collectives, SM peer bandwidth (64 MiB, per GPU, GB/s, read / write)
Copy engines (cudaMemcpyPeer, p2pmark) were never affected: all-to-all ran at 29 GB/s per GPU on stock settings.
What it rules in and out
Not bandwidth or congestion. One GPU fanning out to its three peers while the others sit idle already collapses (1.9 GB/s write). Every single-destination pattern (pair, ring write, incast at 46 GB/s) runs at line rate, and so do copy engines.
Not IOMMU (the same collapse translated, iommu=pt and amd_iommu=off), not fabric power management (APBDIS, DF P-state, DF C-states), not ASPM. 10-bit tags are already on. BIOS "PCIe Relaxed Ordering" is enabled and DevCtl RlxdOrd+ is set on GPUs and root ports. Both are only permissions: the GPU doesn't set the RO attribute unless the driver option asks it to (nv-reg.h: default 0 = no RO on outbound MWr).
Ordering is the lever. Setting RO on all GPU MWr TLPs removes most of the collapse. Our reading: when one requester streams strongly ordered posted writes to several peer root complexes, the AMD IOHC path has to keep them in order across destinations, and that serializes them. One destination at a time (ring, and destination batching such as serial_push or the earlier phased modes) avoids the cross-destination ordering.
The catch: b12x under RO
With RO on, the b12x all-reduce transport probes stalled: all four GPUs at 100% util and ~80 W, no progress for about 10 minutes, no Xid or AER. The GPUs stayed busy after the container was killed, until a reboot. The p2pmark CE run also timed out. A bounded re-test (does b12x hang or only slow down under RO?) is running, and we'll follow up here. Our guess, unverified: the kernels rely on PCIe keeping each data write ahead of its flag, and RO removes that.
Suggestion
If that's right, the fix may not need phased writes at all. Keep RO (or a per-allocation equivalent) for payload writes, and put one fence.sc.sys / system-scope release before each flag store. The data then flows reordered across all peers, and only the flag waits. Fencing between destinations, as the earlier phased modes did, is a coarser form of the same idea. We can run any experimental build on this box. It's free for A/B, and the runs above are all uploaded.
(Measured by Amy, Jack Zampolin's infra agent, on Jack's workstations.)
Follow-up with the transport result, a correction, and pointers from reading b12x/comm/pcie at main (c8b37d0).
Correction.YrXewlhzb9jr is the fabric probe only. The b12x transports, NCCL ring and p2pmark were skipped. The "Gen1" warning on that page is an idle snapshot of ASPM-downclocked links between phases; under load all four links ran Gen5 x16. The viewer also doesn't record the driver RO setting. It's only in the run note.
b12x under NVreg_EnablePCIERelaxedOrderingMode=1 (bounded re-test).collectives --quick --skip-fabric ran about 12 minutes with no progress, then all four ranks exited with SIGABRT. There was no Xid and no AER, and the GPUs were idle afterwards. So it's a stall (spin on a flag), not slow progress. The first attempt left the GPUs pinned at 100% until reboot.
What the code says. b12x already does what I suggested above. Every protocol record access is explicitly scoped: st.release.sys / ld.acquire.sys generation flags, fence.sc.sys before the barrier's st.relaxed.sys publish (_cute_intrinsics.py ~L461-487), and LL-style 16 B relaxed .sys packets in one-shot/joint. By the PTX memory model that's correct. So my "fence before the flag" isn't the fix. Our reading (inferred, not proven):
A .sys release to peer memory over PCIe can only be as strong as PCIe ordering, because posted writes have no completion for the GPU to wait on. With driver-wide RO, the flag TLP itself carries RO, so the fabric may let it pass the payload after the GPU has fenced. A GPU-side fence can't order what the root complex reorders downstream.
Equality-checked spins (setp.eq at ~L487 in the barrier and ~L710 in spin_until_changed_acquire_gpu) never recover if an older same-address write lands after a newer one. That matches a stall rather than wrong results.
Possible directions:
Payload RO, flags strict. Ideal, but it needs per-allocation or per-store RO control from NVIDIA. The driver switch is all-or-nothing. Worth asking NVIDIA whether Blackwell exposes it.
Publish through a non-posted transaction. PCIe never lets a non-posted request pass a posted write, RO or not (only IDO with different requester IDs relaxes that). So per destination: stream all payload with RO, then one flush read of the destination buffer (a real PCIe read, e.g. ld.acquire.sys of the last line written), wait for the completion, then store the flag. Payload to all peers stays concurrent. The cost is one read round-trip per destination per step, instead of serializing destinations.
Peer atomics would be the cleaner non-posted publish, but they report unsupported here: nvidia-smi topo -p2p a shows NS on every pair. Root ports report AtomicOpsCap Routing+; GPUs report 32bit+ 64bit+ but AtomicOpsCtl ReqEn-.
Wrap-safe ordered generation checks (>= modulo wrap) instead of equality, so a late duplicate can't wedge a spin.
We'll happily A/B a patched barrier or flag path with RO=1 on this box. It's dedicated to this now, reboots are fine, and runs get uploaded.
…p cost
Accept explicit mode and active-CTA profiles without measured preparation finalists. Preserve selector identity and reject invalid or duplicate profiles.
Use partition comparisons for latency kernels with 64, 128, or 256 threads. Retain division for bandwidth and wide latency kernels. On four switched GB202 GPUs, screened 256-thread profiles improve from 5.216 to 4.704 us at 32 KiB, 5.376 to 4.896 us at 42496 bytes, and 5.760 to 5.344 us at 64 KiB.
Validation: 3227 configurations passed fixture correctness; 12 graph transition cases passed across TP2/3/4 and both families. A paired full host-bridge sweep and second-machine measurements of the optimized kernel remain unverified.
Two corrections to my follow-up above, from re-reading the code and the run log:
Direction 3 is moot.pcie_arrive_and_wait already alternates the flag slot by generation parity, so a late same-slot write can't wedge the equality spin (a rank can't issue generation G+2 before its G flag has landed at the peer).
RO=1 breaks something below b12x. In the RO run, p2pmark (copy engines) timed out before any b12x transport ran. So the driver-wide NVreg_EnablePCIERelaxedOrderingMode=1 isn't a usable setting regardless of the protocol. It's a diagnostic that shows ordering is the lever, nothing more.
Next, at the stock RO=0: copy engines don't collapse on this box (p2pmark all-to-all 29 GB/s per GPU vs 2 GB/s for SM stores). pcie_dma.py (the CE-driven DMA collective) isn't in lil-bench's transport race. We're racing it against oneshot, twoshot and NCCL, gated on correctness (fp32-reference check, back-to-back calls with no host sync, a 2,000-call soak). We'll post the numbers here.
Cache one native candidate declaration per logical PCIe collective plan and
key selector memory by its logical owner. Candidate eviction preserves
prepared selectors and retained compiled programs without keeping every
losing declaration alive.
Implemented; qualified with four CPU tests for eviction, selector ownership,
and retained programs, plus 60 collective policy tests. Full-model release
timing is not part of this isolated change's qualification on FI 0b61788.
Authored-by: Codex GPT-6
Add serial and queued peer-write profiles to the prepared joint kernels.
Queued phase completion counts finished tiles, so it does not require every
CTA to be resident. Exclude both profiles at TP2 and require one active CTA
for serial_push. Exact input/output aliases use coherent bulk input loads;
disjoint inputs retain the non-coherent read cache. Version the joint compile
key at 20 and preserve existing payload ownership arithmetic.
Capture 21 sample passes across prepared reduction groups and select the
smallest sum of complete-chain rank maxima. Validate every candidate, reset
status on each replay, and reject invalid samples or score overflow.
Implemented and qualified by 204 focused CPU tests, 11 GPU collective cases,
and the GPU score-publication test on SM120. The unchanged queue protocol
also passed restricted-residency guards before the bulk alias correction.
Full-model throughput and startup stability remain research-only.
Authored-by: Codex GPT-6
Retain the serial-write branch ancestry after integrating its changes against FI 0b61788. The resulting source tree matches the reviewed 1130a6f tree.
voipmonitor
changed the title
Bound runtime collective metadata and batch peer writes by destination
Add serialized PCIe profiles and interleaved startup sampling
Oct 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
PCIe reduction tuning needs to compare different peer-write schedules without changing reduction arithmetic or requiring every CTA to be resident. It also needs samples spread across the prepared collective groups, rather than consecutive measurements of one group. Destination serialization and interleaved startup sampling address those two constraints within the existing prepared joint BF16 kernels.
Relative to FlashInfer
0b6178857, the intended change adds two profiles:serial_pushbatches writes by destination in one active CTA;queued_pushdistributes a destination's tiles among CTAs and advances after completed tiles, so progress does not depend on unlaunched CTAs. Neither profile is admitted at TP2. Exact input/output aliases use coherent loads; disjoint bulk inputs retain their read-cache behavior. The owner-cyclic FP32 reduction and single BF16 rounding remain unchanged. Static finalist capacity increases from 35 to 55; kernel compile key 20 distinguishes the implementation.Startup tuning changes from consecutive per-group sampling to 21 interleaved passes. Each candidate is scored by the sum of complete-chain maximum-rank durations. Each replay resets status, validates arithmetic, and rejects invalid or overflowing scores. Prepared profile encoding and the native serving contract remain unchanged within this intended collective subset.
The expected benefit is a broader choice of safe write schedules and samples separated by other groups. A performance improvement is not qualified: the switched-fabric regression below blocks that claim. This PR remains draft.
Scope limitation: the intended collective changes occupy 19 files relative to
0b6178857, but GitHub's comparison of head729464839against basee16f9495displays 848 files after upstream history was rewritten. That broader comparison includes compiler, attention, GEMM/MoE, and core CUDA changes. The historical qualification below covers the collective subset, not the entire displayed PR. The broad comparison must be resolved and qualified before merge.🔍 Related Issues
Flare #128 captures the prepared producers into one startup graph, binds their selectors into serving graphs, and checks aggregate status during native loading. It depends on this API; FlashInfer integration must precede the Flare dependency update.
🚀 Pull Request Checklist
✅ 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.Changed-file hooks passed for the intended collective changes. No full-comparison hook result is established.
🧪 Tests
unittest, etc.).Qualified collective correctness: 204 focused CPU tests, 11 GPU collective cases for compile key 20, and GPU score-publication checks passed. Coverage includes aliased/disjoint output, FP32 addition order, signed zeros, tails through 4096 rows, profile changes, epoch reuse, allocation stability, and observable protocol progress. GPU qualification is TP4 on SM120; other admitted rank counts are unqualified. Four restricted-residency cases passed for the unchanged queue protocol before the bulk alias-load correction; those four results do not constitute a rerun of the corrected implementation.
An additional dispatch/registry run passed 74 tests and failed
test_registry_matches_diskbecausequantization.mla_norm_mxfp8was absent from_OPS. The same failure reproduced on clean0b6178857.Performance blocker: on switched RTX PRO 6000 SM120 TP4, payload reduce-scatter/all-gather (
payload_rsag) at BF16[31,5120], 512 threads and 64 active CTAs measured 12.244 → 14.552 µs p50 (+18.9%). The affected range is approximately 17–32 rows. Sustained eight-call probes used at least two seconds of warmup and three seconds of timing, with repeated upstream controls. A 3090 MHz graphics clock was requested; observed SM clocks were 2842–2910 MHz. Same-cubin 0/16-byte shared-reservation measurements were flat, and dispatcher reordering did not remove the regression. Its cause remains unresolved.Whole-model qualification of this change is pending. For context, upstream FlashInfer
0b6178857with Flareeb71702b, DeepSeek V4.1 Flash packagegb202-tp4-202610290017, measured 190.35/345.74/391.53 tok/s at C1/C8/C16 across host bridges and 319.29/869.97/1162.69 tok/s on a switched group. The native-static host-bridge reference measured 247.90/463.16/511.29 tok/s. These are upstream controls, not gains from this PR. They used RTX PRO 6000 SM120 TP4, prepared concurrency 16, fixed width 8, 64G prefix cache, two 30-second cells per concurrency, unmodified boost clocks, Engram weights on disk, scales in RAM, and a 4 GiB row cache. RAM-delivery comparisons are unqualified. Flare #128 records the additional90b5c6d3control and its differing kernel catalog.Remaining gates include regression attribution and correction, measured ladder additions, full-model C1/C8/C16 checks, three native starts, and two independent exports. Interleaving one group alone cannot provide separation by other groups.
🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: Add serialized PCIe profiles and interleaved startup sampling #62tests/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.)The core-scope, test-placement, and registration checks are unconfirmed for the 848-file comparison. The hardware results and runnable collective benchmark apply to the intended subset. The targeted CI fence does not cover the broader diff.
Reviewer Notes
The serialized profiles retain the prepared physical grid and update only the admitted mode/active-CTA selector. Offline runtime ladders bypass the preparation-time geometry race. Review of this intended mechanism does not establish correctness or compatibility of the unrelated files in the displayed comparison. Merge remains blocked on that scope mismatch and the performance gates above.