Skip to content

fix(tests): initialize SM107 staging kernel capability stub - #6509

Merged
Anerudhan merged 1 commit into
mainfrom
fix/sm107-staging-kernel-fixture
Oct 15, 2026
Merged

Anerudhan merged 1 commit into
mainfrom
fix/sm107-staging-kernel-fixture

Conversation

@YangXu1990uiuc

@YangXu1990uiuc YangXu1990uiuc commented Oct 15, 2026 •

Copy link
Copy Markdown
Collaborator

📌 Description

The portable CUDA DLPack/capture-stream test constructs Sm107BlockScaledSymmBuffer with object.__new__, bypassing its initializer. After #6447, launch() reads self.kernel.masks_invalid_combine_slots to decide whether invalid combine slots need clearing. The fixture omitted kernel, causing all three quantization cases to fail with AttributeError before the dummy compiled callable runs.

Initialize the same lightweight kernel capability stub already used by the host contract fixture. Setting masks_invalid_combine_slots=False preserves the test's existing launch-side clearing path and keeps the real DLPack rebinding and CUDA Graph replay checks portable across supported CUDA GPUs. This is a test-fixture repair; production code is unchanged.

🔍 Related Issues

Fixes the CI regression introduced by #6447 and observed in the A10G job for #6506. The job reported 3 failures, all from this test, alongside 1107 passes.

🚀 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

  • 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.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

  • Tests have been added or updated as needed.

  • All tests are passing (unittest, etc.).

  • Changed-file pre-commit hooks and git diff --check: passed.

  • python -m pytest tests/moe_ep/test_sm107_block_scaled_contracts.py tests/moe_ep/test_sm107_torch_staging.py -q: 117 passed, 11 skipped on the CPU-only local host. All 11 CUDA staging cases skipped because CUDA is unavailable locally; this is not GPU validation.

  • Normal GPU CI was started with the documented run-ci label at head 38ae833e6. It includes the three originally failing real-DLPack/capture-stream cases; results are pending.

🔬 Experimental Track

  • This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
    • 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.
# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

The production initializer always assigns self.kernel; the missing attribute is specific to this manually constructed test workspace. The corresponding host fixture was updated in #6447, but this CUDA fixture was missed. The original PR's targeted CI scope also omitted test_sm107_torch_staging.py.

Self-review found no remaining changes needed. The all-files hook checkbox remains unchecked because hooks ran only on the changed file.

Summary by CodeRabbit

  • Tests
    • Updated coverage for DLPack capture-stream launches to confirm invalid combine slots continue to be cleared. This is a test-only change and does not alter application behavior or add user-facing functionality.

The portable DLPack capture test bypasses workspace initialization. Provide the kernel capability now read by launch(), matching the host contract fixture and retaining launch-side combine clearing.

AI-assisted change.
@coderabbitai

coderabbitai Bot commented Oct 15, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: defaults
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 7daf756b-a4b9-478b-afe0-3dba98d9226c
📥 Commits

Reviewing files that changed from the base of the PR and between 021f3f5 and 38ae833.

📒 Files selected for processing: 1
  • tests/moe_ep/test_sm107_torch_staging.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The DLPack capture-stream test’s dummy kernel now sets masks_invalid_combine_slots=False. A comment states that launch-side clearing of invalid combine slots remains enabled.

Changes: DLPack capture-stream test

Layer / File(s) Summary
Declare dummy kernel masking behavior
tests/moe_ep/test_sm107_torch_staging.py
The dummy kernel now sets masks_invalid_combine_slots=False. A comment notes that launch-side clearing remains enabled.

Priority: ⬇️ Low

Estimated code review effort: 1 (Trivial) | ~2 minutes

Change: Other

Merge Risk: ⚪ Minimal · up to 38ae8

This test-fixture repair has no identified merge-blocking issue. GPU CI remains the appropriate validation for the CUDA cases.

Pre-merge checks: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 4 passed
Check name Status Explanation
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Title check Passed The title clearly identifies the test-only fix: initializing the SM107 staging kernel capability stub.
Description check Passed The description explains the failure, the fixture change, the rationale, test results, CI limitations, and scope. It follows the repository template and clearly notes that GPU validation is pending an…
  • Fix all pre-merge checks with AI
✨ Finishing Touches: 💡 2
📝 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@YangXu1990uiuc YangXu1990uiuc left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex bot review · model: gpt-6-astra

Reviewed 38ae833e6d93 against 021f3f574581. No blocker found; this is an appropriate fixture repair.

The production constructor assigns self.kernel unconditionally after building it. This test deliberately bypasses that constructor and supplies a dummy compiled callable, so the missing attribute belongs in the fixture. Setting masks_invalid_combine_slots=False retains the real launch-side clearing path; it does not disable the DLPack rebinding or capture/replay checks, and production code is unchanged.

Independently verified on H100 NVL, with Torch 2.14.1 and CuTeDSL 4.8.0:

  • Before: restored the base test file while keeping identical production code; all three quantization cases failed with the reported missing-kernel AttributeError.
  • After: the two complete files test_sm107_block_scaled_contracts.py and test_sm107_torch_staging.py passed: 128 passed, zero skipped, including all 11 portable CUDA cases and the three real-DLPack/new-capture-stream cases. Source paths and hashes were verified.

This validates the portable staging/launch contract with the existing dummy callable; it does not requalify the native SM107 kernels. The normal GPU CI is still pending. Posting COMMENT because this is the same GitHub account as the author.

@Anerudhan
Anerudhan enabled auto-merge (squash) October 15, 2026 03:19
@Anerudhan
Anerudhan merged commit 3637083 into main Oct 15, 2026
29 of 30 checks passed
@Anerudhan
Anerudhan deleted the fix/sm107-staging-kernel-fixture branch October 15, 2026 04:11
yihuawei added a commit to yihuawei/flashinfer that referenced this pull request Oct 15, 2026
…merge

The fork's main took upstream through flashinfer-ai#6509 via clean merges; those commits
are already part of the e43f75d merge, so the tree is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Anerudhan added a commit that referenced this pull request Oct 15, 2026
<!-- .github/pull_request_template.md -->

## 📌 Description

`auto` restricts cuDNN to a small GB200 region and returns to slower
native kernels at 32K/64K. Wide FP32 calls also miss the persistent
backend.

Use coarse head/length regions on SM100/SM103, favoring cuDNN when its
GPU time is within 10% of the fastest valid measured backend, with an
explicit exception for GB300 wide-head BF16 calls. Retain native
math/layout checks, explicit backend overrides, and
state-pool/checkpoint handling. A supplied native workspace no longer
blocks cuDNN; cuDNN manages its own scratch, and native fallback uses
the supplied workspace.

- Local heads: few 1–16; moderate 17–64; wide 65–96.
- Sequence length: short 1K–<4K; medium 4K–<8K; long 8K–64K. K = 1024.
- FP32 wide-head calls use persistent. BF16 wide-head calls use cuDNN on
both GB200 and GB300. GB300 accepts the measured 14–15% GPU overhead
against Cake.
- Few-head FP32 calls switch to cuDNN at 4K; BF16 switches at 4K for up
to 8 heads and 8K for 9–16 heads.
- Moderate-head BF16 calls keep small-bh. Moderate-head FP32 calls use
persistent below 4K and cuDNN above it, except the GB300 8K–<16K bucket.

The first-run data covers 96 valid bounded-gate cases at H=8/12/64/96
and T=1K/4K/8K/16K/32K/64K. The measured `c29969ef` policy selected
cuDNN in 53 cases; the largest measured cuDNN overhead against the
fastest valid backend is 8.04%. The 10% limit is an offline tuning
target; GB300 wide-head BF16 is now an explicit exception.

Measured baseline → `c29969ef` auto GPU timings, in µs/call (BF16
activations, D128, one packed sequence):

| GPU | Heads | State | Tokens | Baseline auto | PR auto | PR backend |
|---|---:|---|---:|---:|---:|---|
| GB200 | 8 | FP32 | 65536 | 1950.70 | 561.97 | cudnn |
| GB300 | 8 | FP32 | 65536 | 1940.20 | 561.01 | cudnn |
| GB200 | 64 | FP32 | 4096 | 207.10 | 200.48 | cudnn |
| GB300 | 64 | FP32 | 4096 | 203.78 | 195.91 | cudnn |
| GB200 | 96 | BF16 | 8192 | 450.85 | 445.78 | cudnn |
| GB300 | 96 | BF16 | 8192 | 457.43 | 395.72 | cake |
| GB200 | 96 | FP32 | 8192 | 452.29 | 358.67 | cute-dsl-persistent |
| GB300 | 96 | FP32 | 8192 | 458.95 | 355.27 | cute-dsl-persistent |

Paired same-node runs compare the requested baseline `fd14b3056d0c` with
PR commit `c29969ef04cb`. The comparison includes intervening
main-branch changes. Each run uses 5 warmups, 30 samples, reused inputs,
state reset outside timing, and no L2 flush. GPU time is the median
summed eager CUPTI kernel duration.

| GPU | Valid paired cases | CPU speedup | GPU speedup | Graph speedup |
|---|---:|---:|---:|---:|
| GB200 | 90 | 0.887× | 1.117× | 1.190× |
| GB300 | 90 | 0.881× | 1.183× | 1.203× |

Speedup is the geometric mean of baseline/PR ratios; below 1 means
slower. CPU launch time increases by about 13% on GB200 and 14% on
GB300. This policy prioritizes GPU time.

The `c29969ef` runs selected cuDNN in 53 bounded-gate sequence cases.
The largest GPU penalty against the fastest valid explicit backend is
**9.05% on GB200** and **7.80% on GB300**; no cuDNN choice exceeds 10%.
All 53 choices also use cuDNN during graph capture with a supplied
native workspace.

[Full results, raw data, tests, and GPU
chart](https://gitlab-master.nvidia.com/agopal/flashinfer_kda_bench/-/blob/main/results/2026-10-15_c29969ef04cb_feat-kda-cudnn-auto-policy/RESULTS.md)
(NVIDIA access).

The public-call comparison can also be reproduced with the existing
benchmark, which reports host/wall and graph timing rather than summed
eager CUPTI time:

```bash
python benchmarks/bench_cudnn_linear_attention.py --family kda --lengths 65536 --heads 8 --state-dtype fp32 --backends auto,cudnn,cute-dsl --output kda-h8-t64k.json
python benchmarks/bench_cudnn_linear_attention.py --family kda --lengths 8192 --heads 96 --state-dtype fp32 --backends auto,cudnn,cute-dsl-persistent --output kda-h96-t8k.json
```

## 🔍 Related Issues

Extends the cuDNN KDA policy from #6187. Independent of the native
launch optimization in #6297.

## 🚀 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

- [x] 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.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All targeted tests are passing (`unittest`, etc.).

- Latest GB300 wide-head BF16 policy change: `pytest
tests/kda/test_cudnn_auto_dispatch.py -q` — **113 passed, 26 skipped**.
Local cuDNN frontend 1.29.0 does not meet the GPU tests’ 1.31
requirement. Changed-file pre-commit checks passed.
- Earlier policy: **133 passed** on a local SM100 GPU (148 SMs), PyTorch
2.13.0+cu130, CuTe DSL 4.7.0, cuDNN frontend 1.31.0.
- Covers both architecture policies on the host, real backend dispatch,
output and final-state checks, CUDA Graph replay, persistent alignment
and SM-count guards, optional runtime fallback, and explicit backend
overrides.
- Changed-file pre-commit checks passed, including mypy, Ruff
lint/format, whitespace and conflict checks. The full repository suite
and `pre-commit --all-files` were not run.
- Measured commit `c29969ef` passed **133 tests on GB200 and 133 on
GB300**, including cuDNN capture with a supplied native workspace and
native fallback capture. Lyris jobs 3353636/3353637 completed all 1,728
PR rows plus 216 paired baseline-auto rows. Statuses match the first-run
matrix; known Kimi-Linear failures remain visible. Cluster runtime:
PyTorch 2.13.0+cu132, CuTe DSL 4.7.1, cuDNN frontend 1.31.0.
- The earlier CI run failed in three SM107 MoE staging tests because the
fixture lacked `kernel`. Cherry-picked the test-only fix from #6509
(`38ae833e6`) as `91b36d76`; KDA code and benchmark results are
unchanged.
- `pytest tests/moe_ep/test_sm107_block_scaled_contracts.py
tests/moe_ep/test_sm107_torch_staging.py -q`: **128 passed** on the
local SM100 GPU, including all three previously failing CUDA capture
cases. Changed-file pre-commit checks passed.
- [Full GitHub
CI](https://github.com/flashinfer-ai/flashinfer/actions/runs/38679413398)
was restarted with `run-ci` on `91b36d76`; results are pending.

## 🔬 Experimental Track

<!-- Only for PRs submitted under the experimental policy
(CONTRIBUTING.md → "Experimental APIs and Backends").
     Leave this section untouched for normal PRs. -->

- [ ] This PR is **experimental**: it adds or changes code under
`flashinfer/experimental/` and/or an `@flashinfer_experimental_api`.
Tracking issue: #
- [ ] 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.

<!-- Required for experimental PRs. Replace the commented lines below
with your targets.
Do not delete the fence or change its `experimental-tests` tag — the
experimental-track
watcher reads it verbatim to decide which targets to ask CI for. -->

```experimental-tests
# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.
```

## Reviewer Notes

- GB300 BF16 state with 65–96 heads now selects cuDNN for 1K–64K tokens.
Cake remains the native fallback. Historical auto timing tables above
describe `c29969ef` before this change; explicit cuDNN measurements show
14–15% overhead against Cake. FP32 selection is unchanged.

- PTX stays explicit: its state-value bound is narrower than native auto
semantics. For GB300 moderate-head FP32 at 8K–<16K, choose persistent
because cuDNN was 16.4% slower than PTX at 8K. Persistent itself was
10.1% slower than PTX at that point.
- Auto persistent selection is restricted to one sequence with H no
greater than the SM count, so no inter-CTA handoffs or split-grid
concurrency restriction is introduced.
- Auto-selected cuDNN leaves a supplied native `prefill_workspace`
unused and owns its capture scratch. Native fallback retains the
supplied workspace. Tests cover cuDNN selection and native fallback
during capture; warm with the same arguments and stream before capture.
- No kernel math, public signature, decode dispatch, or experimental
backend was changed.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Automatic prefill selection now covers eligible single-sequence
workloads on SM100/SM103 with 128-dimensional heads, up to 96 heads, and
sequence lengths from 1K to 64K tokens. Selection accounts for state
type, head count, device, and cuDNN availability.
* If cuDNN is unavailable or declines an eligible workload, selection
can use a persistent native backend when its eligibility checks pass;
otherwise, it uses another applicable native backend.

* **Documentation**
* Clarified backend eligibility and fallback behavior, and how workspace
handling differs between cuDNN and native calls.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Yang Xu <yanxu@nvidia.com>
jimmyzho pushed a commit that referenced this pull request Oct 20, 2026
<!-- .github/pull_request_template.md -->

## 📌 Description

The portable CUDA DLPack/capture-stream test constructs
`Sm107BlockScaledSymmBuffer` with `object.__new__`, bypassing its
initializer. After #6447, `launch()` reads
`self.kernel.masks_invalid_combine_slots` to decide whether invalid
combine slots need clearing. The fixture omitted `kernel`, causing all
three quantization cases to fail with `AttributeError` before the dummy
compiled callable runs.

Initialize the same lightweight kernel capability stub already used by
the host contract fixture. Setting `masks_invalid_combine_slots=False`
preserves the test's existing launch-side clearing path and keeps the
real DLPack rebinding and CUDA Graph replay checks portable across
supported CUDA GPUs. This is a test-fixture repair; production code is
unchanged.

## 🔍 Related Issues

Fixes the CI regression introduced by #6447 and observed in [the A10G
job for
#6506](https://github.com/flashinfer-ai/flashinfer/actions/runs/38667451074/job/116202702238?pr=6506).
The job reported 3 failures, all from this test, alongside 1107 passes.

## 🚀 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

- [x] 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.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

- Changed-file pre-commit hooks and `git diff --check`: passed.
- `python -m pytest tests/moe_ep/test_sm107_block_scaled_contracts.py
tests/moe_ep/test_sm107_torch_staging.py -q`: **117 passed, 11 skipped**
on the CPU-only local host. All 11 CUDA staging cases skipped because
CUDA is unavailable locally; this is not GPU validation.
- [Normal GPU
CI](https://github.com/flashinfer-ai/flashinfer/actions/runs/38672336957)
was started with the documented `run-ci` label at head `38ae833e6`. It
includes the three originally failing real-DLPack/capture-stream cases;
results are pending.

## 🔬 Experimental Track

<!-- Only for PRs submitted under the experimental policy
(CONTRIBUTING.md → "Experimental APIs and Backends").
     Leave this section untouched for normal PRs. -->

- [ ] This PR is **experimental**: it adds or changes code under
`flashinfer/experimental/` and/or an `@flashinfer_experimental_api`.
Tracking issue: #
- [ ] 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.

<!-- Required for experimental PRs. Replace the commented lines below
with your targets.
Do not delete the fence or change its `experimental-tests` tag — the
experimental-track
watcher reads it verbatim to decide which targets to ask CI for. -->

```experimental-tests
# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.
```

## Reviewer Notes

The production initializer always assigns `self.kernel`; the missing
attribute is specific to this manually constructed test workspace. The
corresponding host fixture was updated in #6447, but this CUDA fixture
was missed. The original PR's targeted CI scope also omitted
`test_sm107_torch_staging.py`.

Self-review found no remaining changes needed. The all-files hook
checkbox remains unchecked because hooks ran only on the changed file.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Tests**
* Updated coverage for DLPack capture-stream launches to confirm invalid
combine slots continue to be cleared. This is a test-only change and
does not alter application behavior or add user-facing functionality.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

(cherry picked from commit 3637083)
triple-mu added a commit that referenced this pull request Oct 24, 2026
<!-- .github/pull_request_template.md -->

## 📌 Description

Adds `UlyssesCommunicator(backend="pcie")`: an experimental single-node
all-to-all transport for PCIe-connected GPU groups without NVLink (e.g.
8x RTX PRO 5000), targeting Ulysses sequence parallelism in video/image
diffusion serving. Being experimental, it is reachable from
`backend="auto"` only where NVLink is unavailable and
`FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1` is set identically on
every rank; naming it explicitly never needs that variable. Either way
it warns once.

**Why.** The existing NVLink backend lets SM threads store into
peer-visible staging buffers; across a PCIe host bridge that degenerates
into small transactions (measured 5-13x slower than copy engines during
bring-up, see `docs/design_docs/ulysses_pcie.md`). NCCL works but leaves
bandwidth and stability on the table on this fabric (see numbers below).

**What.**

- Routes: world size 1 is an identity path; 2 uses all-pairs CUDA P2P
copy engines (`cudaMemcpy2DAsync` writes land directly in the peer's
final layout, no pack/unpack); 4/8 prefer an all-RDMA route (per-rank
mlx5 RC QP, interleaved-UMR MKeys, GPUDirect RDMA) and fall back to
all-P2P with a `RuntimeWarning`; an eight-rank 4+4 NUMA hybrid
(same-NUMA P2P, cross-NUMA mlx5) is available via
`FLASHINFER_ULYSSES_PCIE_ROUTE`.
- Transforms: `scatter_heads` / `gather_heads` interleave the head axis;
`exchange_chunks` takes `[1, 1, world_size, chunk]` to the same shape
(`all_to_all_single` over a payload a preceding pack already laid out
destination-major) and interleaves nothing. On the RDMA routes it
configures no MKey, which is what frees it from the 16-bit `bytes_count`
that caps the head-axis transforms at 65535 B per peer.
- Synchronization: full-group epoch barriers in device memory
(release/acquire, sticky abort half for the RDMA routes); the all-P2P
route is CUDA-graph capturable, the RDMA routes refuse capture and bound
NIC stalls with a 10 s deadline plus a group-wide abort/quiesce
protocol.
- Lifecycle: explicit `allocate_output` registration (collective, lives
until `close()`); `input_buffer()` exposes the RDMA landing buffer so
producers can write operands where the NIC reads them, removing the
largest device-to-device copy on the path (measured ~6% per call at
layer scale). The backend has no staging workspaces: `workspace=`,
`create_workspace` and the head-chunk primitives from #5027 are refused
at world size above one with a message that names `allocate_output()`.
- Failure envelope: a rejection before native code runs (bad operand,
refused capture, second stream) is decided from the call's own
arguments, so under the SPMD contract every rank rejects it -- the RDMA
routes raise the ordinary error and stay usable, the all-P2P route (no
abort protocol) poisons teardown and enters BROKEN. A failure inside the
transport is fail-stop on every route: on the RDMA routes the failing
rank publishes the sticky abort, the peers leave the barrier and
`close()` completes; on all-P2P `close()` refuses to synchronize and the
process must exit. The design note has the table.
- The element type is per-call (`dtype=` on the collectives), decoupled
from the communicator: the transport moves opaque bytes, so FP8/INT8
payloads ride a BF16-sized communicator within the same byte budget. The
constructor's capacity is therefore `max_bytes`, replacing main's
element-count `max_elems`: the other communication workspaces in
`flashinfer.comm` are sized in bytes as well (`buffer_size_in_bytes`,
`workspace_size`), and an element count has no single meaning once the
element type varies per call. #4924 already targets `max_bytes`. The
staging-buffer API from #5027 (`UlyssesWorkspace(max_elems=...)`,
`create_workspace(max_elems=...)`) keeps its element count, since those
buffers hold one dtype.
- Docs: design note (`docs/design_docs/ulysses_pcie.md`),
`docs/api/comm.rst`, env-var reference; fi_trace templates for all three
transforms (`scatter_heads` / `gather_heads` gain `trace=` decorators).
- Housekeeping: six `flashinfer.comm` re-exports of JIT module helpers
with no users in the repository (`gen_trtllm_comm_module`,
`gen_vllm_comm_module`, `gen_pcie_ipc_comm_module`,
`get_pcie_ipc_comm_module`, `gen_ulysses_a2a_module`,
`get_ulysses_a2a_module`) are dropped; the helpers stay importable from
their own modules and from `flashinfer.jit`.

**Changes since the approved revision (a481f1d).** The branch was
rebased onto main (#5027 workspace API, #5263 generated NVLink kernels)
and the following landed on top, in this order:

1. `fix(comm): finish the max_bytes rename in the head-chunk tests,
benchmark and export example` -- four `max_elems=` constructor calls
that arrived with #5027 / #5263 (`_destination_passing_body`,
`_head_chunk_body`, `bench_ulysses_head_chunk.py`,
`examples/pytorch/ulysses_a2a_export/interface.py`) raised `TypeError`
on any 2+ GPU host.
2. `format: format code style by pre-commit` -- the two hooks CI
flagged.
3. `test(comm): share the Ulysses test helpers and gloo fixture` -- one
FileStore `gloo_pg` fixture in `tests/comm/conftest.py`, `full_mesh` /
`forbid_ipc_and_jit` / rendezvous helpers in
`tests/test_helpers/ulysses.py`.
4. `fix(comm): check the group's CUDA all-to-all support before arming
NVLink` -- the check ran only for `backend="nccl"`, after the NVLink IPC
exports were armed.
5. `fix(comm): reject cross-rank FLASHINFER_ULYSSES_PCIE_NICS
disagreement` -- the NIC override is now compared across ranks like the
route and GID overrides.
6. `feat(comm): warn at garbage collection when a Ulysses communicator
was not closed` -- `ResourceWarning` from `__del__` for an armed
communicator; no teardown from the finalizer.
7. `fix(comm): refuse workspaces and the head-chunk primitives on the
multi-rank PCIe backend` -- see Lifecycle above.
8. `fix(comm): publish the abort when the RDMA routes reject an operand
natively` -- the RDMA routes' per-call operand checks run inside the
native abort envelope, so a rank whose operand native rejects releases
its peers instead of leaving them in the opening barrier (the only
native change: the checks moved into `ValidateExchangeOperands`,
protocol and wire format untouched); the caller's stream is bound only
after a call passes validation; the Python-side envelope is unchanged
and documented.
9. `refactor(comm): route Ulysses output validation through the shared
helpers` -- `_output_geometry` is the single source of the output shape;
the overlap test reuses `_storage_ranges_overlap`.
10. `docs(comm): update Ulysses PCIe prose for auto selection, the
failure envelope and follow-ups`.
11. Rebased onto main 28ae778 (2026-09-24). The only conflict was
`flashinfer/trace/templates/comm.py`, where #4929 and this PR both
append templates at the end of the file; both blocks are kept, in main's
order, and no other file changed in the rebase.
12. Merged main `7bcf91831` (2026-09-29): three additive conflicts
(`flashinfer/aot.py`, `flashinfer/comm/__init__.py`,
`tests/trace/template_registry.py`), both sides kept. Followed by `Find
Conda rdma-core build inputs for PCIe Ulysses JIT` (72c958a):
`gen_ulysses_pcie_module` adds `<sys.prefix>/include` and
`-L<sys.prefix>/lib` when `infiniband/mlx5dv.h` / `libmlx5.so` live
there (the CI Conda environment), and
`missing_ulysses_pcie_dependencies()` now also probes for the
`infiniband/verbs.h` / `infiniband/mlx5dv.h` headers (one host-mock test
added; design note updated).
13. Merged main `7aa42e2dc` (2026-10-04). Two additive conflicts:
`scripts/task_test_multi_gpu_comm_kernels.sh` (both sides extend
`SPAWN_MANAGED_TEST_FILES`; the union is kept) and
`tests/trace/example.py` (main sorted the `flashinfer.*` imports;
`import flashinfer.comm` takes its place in that order). `CLAUDE.md`,
`docs/api/comm.rst`, `flashinfer/aot.py` and
`tests/trace/test_fi_trace_template_consistency.py` merged
automatically. No file under `flashinfer/comm/ulysses*`,
`flashinfer/jit/comm.py`, `csrc/ulysses*`,
`include/flashinfer/comm/ulysses_pcie.cuh` or `tests/comm/test_ulysses*`
changed on main between the two bases.
14. Merged main `188bdd769` (2026-10-05). Two modify/delete conflicts:
#5953 removed `examples/pytorch/ulysses_a2a_export/` and moved its
fixture to `tests/comm/cake_ulysses_a2a_fixture.py`, so the deletion is
taken and the `max_bytes` rename from item 1 is re-applied to the moved
fixture in `fix(comm): finish the max_bytes rename in the moved cake
Ulysses fixture` (57ed509); `docs/api/comm.rst` and `flashinfer/aot.py`
merged automatically. No file under `flashinfer/comm/ulysses*`,
`flashinfer/jit/comm.py`, `csrc/ulysses*`,
`include/flashinfer/comm/ulysses_pcie.cuh` or `tests/comm/test_ulysses*`
changed on main between the two bases.
15. Rebased onto main `9c6797eda` (2026-10-05) at the maintainers'
convenience for a linear history: the 28 non-merge commits above replay
in order and the five merge commits (1f82831, 25bb39c, f7ce2ba,
5e8919a and their resolutions) are gone; `8668d921` is the new head.
The replay re-resolved the same conflicts
(`flashinfer/comm/__init__.py`, `flashinfer/aot.py`,
`tests/trace/example.py`, `tests/trace/template_registry.py`,
`scripts/task_test_multi_gpu_comm_kernels.sh`, the removed example
files), and the result differs from the merged head `57ed5097` in
exactly one place: `flashinfer/comm/__init__.py` no longer re-exports
`gen_pcie_ipc_comm_module` / `get_pcie_ipc_comm_module`, which the merge
in item 12 had brought back by accident; this is the drop described
under Housekeeping, and the repository has no caller of either name
through `flashinfer.comm`.
16. Rebased onto main `67c190f7d` (2026-10-06); `0f60c090` is the new
head, still 28 linear commits. One conflict: the expected-output list in
`tests/trace/example.py`'s docstring, where #6035 appended the
`moe_layer_cudnn_frost_*` entries at the spot this PR appends the
`ulysses_*` entries; both lists are kept, in main's order. `CLAUDE.md`,
`flashinfer/aot.py` and `tests/trace/template_registry.py` merged
automatically. No file under `flashinfer/comm/ulysses*`,
`flashinfer/jit/comm.py`, `csrc/ulysses*`,
`include/flashinfer/comm/ulysses_pcie.cuh` or `tests/comm/test_ulysses*`
changed on main between the two bases.
17. Rebased onto main `87ef73934` (2026-10-14); `e7c7bbd50` is the new
head, 29 linear commits. One conflict: the tail of
`tests/trace/example.py`, where #5304 appended its `top_k_varlen`
example calls at the spot this PR appends the single-rank Ulysses block;
both are kept, in main's order (the docstring's expected-output list
merged automatically this time). `CLAUDE.md`, `docs/api/comm.rst`,
`flashinfer/aot.py`, `flashinfer/comm/__init__.py`,
`scripts/task_test_multi_gpu_comm_kernels.sh` and
`tests/trace/template_registry.py` merged automatically. #6332
(`flashinfer/comm/qkv_norm_rope_exchange_overlap.py` and
`benchmarks/comm/bench_qkv_norm_rope_exchange_overlap.py`) arrived on
main with three `UlyssesCommunicator(..., max_elems=...)` calls, so the
29th commit, `fix(comm): finish the max_bytes rename in the QKV
norm/RoPE exchange overlap`, passes `max_bytes` (the same bf16 element
counts times `torch.bfloat16.itemsize`) exactly as item 14 did for the
moved fixture; nothing else on main constructs a communicator with the
old keyword. No file under `flashinfer/comm/ulysses*`,
`flashinfer/jit/comm.py`, `csrc/ulysses*`,
`include/flashinfer/comm/ulysses_pcie.cuh`, `tests/comm/test_ulysses*`
or `flashinfer/trace/templates/comm.py` changed on main between the two
bases.
18. Rebased onto main `3637083a3` (2026-10-15); `f0ff368c1` is the new
head, still 29 linear commits, no conflicts: #5701 extends
`TORCHRUN_TEST_FILES` in `scripts/task_test_multi_gpu_comm_kernels.sh`
next to the `SPAWN_MANAGED_TEST_FILES` line this PR extends, and both
land. The base now includes #6509, the fix for the
`tests/moe_ep/test_sm107_torch_staging.py` failure that was the only
failing lane in this PR's 2026-10-15 CI run (unrelated to this PR, which
touches no `moe_ep` file). No file under `flashinfer/comm/ulysses*`,
`flashinfer/jit/comm.py`, `csrc/ulysses*`,
`include/flashinfer/comm/ulysses_pcie.cuh`, `tests/comm/test_ulysses*`
or `flashinfer/trace/templates/comm.py` changed on main between the two
bases.

**Performance** (8x RTX PRO 5000, dual-NUMA 4+4, PCIe Gen5 x16, one 400G
mlx5 per GPU; `benchmarks/comm/bench_ulysses_pcie.py`, seq 37888 / heads
56 / head_dim 128 / bf16; median of host-clock loops, max-reduced across
ranks, vs the `all_to_all_single`+permute reference; PCIe-side sd ≤ 8%,
mostly < 1%, while the NCCL reference jitters up to 50-104% at eight
ranks):

| world | route (auto's pick in bold) | scatter | gather | qkv layer |
per-rank remote BW |
|---|---|---|---|---|---|
| 2 | **all-P2P** | 2.02x | 2.00x | 2.01x | 50.5 GB/s |
| 4 | **all-RDMA** | 1.97x | 1.97x | 1.97x | 42.3 GB/s |
| 8 | **all-RDMA** | 1.49x | 1.43x | 1.48x | 42.6 GB/s |
| 8 | 4+4 hybrid | — | — | 1.27x | 36.4 GB/s |
| 8 | all-P2P (correctness fallback) | 0.81x | 0.43x | 0.30x | 8.8 GB/s
|

The transport moves opaque bytes, so the element type should change how
long a call
takes but not how fast the link runs. Same node, ws8 all-RDMA, seq 37888
/ 56 heads /
head_dim 128, only the dtype varies:

| dtype | bytes/elem | scatter ms | qkv layer ms | remote BW (scatter) |
|---|---|---|---|---|
| bf16 | 2 | 1.389 | 5.583 | 42.8 GB/s |
| fp8_e4m3fn | 1 | 0.695 | 2.871 | 42.7 GB/s |
| int8 | 1 | 0.696 | 2.870 | 42.7 GB/s |

Halving the element size halves the time (1.389 -> 0.695 ms) and leaves
bandwidth put;
the two 1-byte types are within noise of each other. ws2 P2P saturates
~95% of the
effective Gen5 x16 link; the RDMA route reaches ~85% of the 400G NIC.

Re-measured at this head on the same node (ws8 all-RDMA, same shapes, 5
trials x 200 iterations, other tenants idle) to confirm the round
changed nothing on the hot path: scatter 1.427 ms at 41.6 GB/s remote
bandwidth (table above: 1.389 ms / 42.8 GB/s), gather 1.554 ms at 38.2
GB/s, both with trial-to-trial sd below 2%; the q/k/v layer loop read
6.93 ms with sd 9.2%, which is inside its own noise band on a shared
node and is not re-plotted here. The `rnr_nak_retry_err` counters on all
eight NICs did not move during the run.

**Limitations.** Experimental; single host; world sizes 1/2/4/8 (the
copy scheduler's `rank ^ step` order needs powers of two); the RDMA
routes require batch 1 and refuse CUDA-graph capture, and bound the
head-axis transforms at head-row pitch `H*D*elem <= 65535` B
(`exchange_chunks` registers no MKey and is not bounded); the JIT module
links libibverbs/libmlx5 even when topology picks all-P2P, so it is
JIT-only and deliberately absent from AOT; a transport failure is
fail-stop as described above, and an all-P2P failure leaves the process
to be terminated by the launcher. Validated on one topology (4+4 NUMA,
one mlx5 per GPU); no work-request-level mlx5 unit tests yet — the RDMA
data path is covered end to end. No CI lane exercises the PCIe wire path
(the CI images carry no rdma-core and the GPU lanes are single-GPU); the
host-mock and single-rank halves of the suite run there, the multi-rank
half runs on the reference node.

Planned follow-ups, prepared on a branch and deliberately kept out of
this PR to leave the approved protocol untouched: publish the abort from
Python on every route (the all-P2P route already imports every peer's
signal buffer), so a Python-side rejection on a rank that broke the SPMD
contract releases the peers too; publish a geometry word with the
opening epoch so ranks that call with different shapes at the same
position break the group instead of writing a misdescribed layout; post
the RDMA receive before the opening barrier (closes the window in which
a peer's write can arrive before the receive is posted; no RNR retries
were counted during this round's benchmark, so this is hygiene rather
than a measured stall); build the all-P2P route without rdma-core;
`out=None` with lazy registration; per-output `capacity_bytes`.

## 🔍 Related Issues

N/A

## 🚀 Pull Request Checklist

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

Ran on the 8-GPU reference node (8x RTX PRO 5000, dual-NUMA 4+4, one
mlx5 per GPU, rdma-core 63, CUDA 13.4, torch 2.13) with all eight GPUs
at `ef9a91a22`, the head before the rebase (2026-09-20):
`tests/comm/test_ulysses_communicator.py` 132 passed / 22 skipped / 0
failed, `tests/comm/test_ulysses_topology.py` 112 passed / 0 skipped
(244 passed / 22 skipped in one run),
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed. The 22 skips are the
forced-NVLink family on this NVLink-less machine; every PCIe case ran
(p2p, all-RDMA and 4+4 hybrid routes at world sizes 2/4/8, the chunk
exchange, the landing-buffer producer path, the p2p fail-stop pair, and
the two-rank RDMA case in which rank 0's native dtype rejection
publishes the abort, rank 1 leaves the barrier with "peer 0 aborted the
exchange", and both ranks close). Correctness is bit-exact against the
`all_gather`-based reference for fp16/bf16/fp32 and byte-exact for
int8/uint8/fp8 across all three routes, including non-power-of-two head
dims on the UMR path, CUDA-graph capture+replay of a q/k/v/out layer on
the P2P route, and `backend="auto"` with the experimental opt-in
selecting PCIe on this node.

Re-run after the rebase at the pushed head `4fec510aa` (2026-09-24) on
the same node, which was shared with another tenant on GPUs 0-3, so the
suite ran on GPUs 4-7 (`CUDA_VISIBLE_DEVICES=4,5,6,7`):
`tests/comm/test_ulysses_communicator.py` 122 passed / 32 skipped / 0
failed, `tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed. The skips are the
forced-NVLink family plus the world-size-8 cases (all-RDMA and the 4+4
hybrid at eight ranks), which need eight visible GPUs; every world-size
2/4 case ran (p2p and all-RDMA routes, the chunk exchange, the
landing-buffer producer path, the fail-stop pairs). The eight-rank paths
stand on the run above: no file under `flashinfer/comm/ulysses*`,
`csrc/ulysses*` or `tests/comm/test_ulysses*` changed on main between
the two bases, and the rebase touched only the trace template file.

Re-run after the 2026-10-04 merge at `25bb39c4` on the same node, first
on GPUs 4-7 while GPUs 0-3 were held by another tenant
(`tests/comm/test_ulysses_communicator.py` 123 passed / 32 skipped),
then on all eight GPUs once they freed up:
`tests/comm/test_ulysses_communicator.py` 133 passed / 22 skipped / 0
failed, `tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed; `pre-commit run
--from-ref origin/main --to-ref HEAD` is green,
`scripts/check_pr_document.py` reports no findings and
`scripts/check_pr_api_diff.py` only the `max_bytes` advisory. The 22
skips are the forced-NVLink family; every PCIe case ran, including the
eight-rank all-RDMA and 4+4 hybrid routes. The head `f7ce2baf` on top
merges main `4b6d64be` (#6040, `cake_selective_state_update` only) and
touches no Ulysses file.

Re-run after the 2026-10-05 merge at the pushed head `57ed5097` on the
same node with all eight GPUs: `tests/comm/test_ulysses_communicator.py`
133 passed / 22 skipped / 0 failed,
`tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed, and
`tests/comm/test_cake_ulysses_a2a.py` + `tests/comm/test_ulysses_a2a.py`
+ `tests/comm/test_ulysses_head_chunk.py` 16 passed / 4 skipped (the
generated NVLink portfolio skips on this SM120 node; the moved fixture
was additionally exercised by constructing its communicator with
`max_bytes` in a one-rank process). The skips are the forced-NVLink
family. `pre-commit run --from-ref <main> --to-ref HEAD` is green.

Re-run at the rebased head `8668d921` on the same node with all eight
GPUs: `tests/comm/test_ulysses_communicator.py` 133 passed / 22 skipped
/ 0 failed, `tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed,
`tests/comm/test_cake_ulysses_a2a.py` + `tests/comm/test_ulysses_a2a.py`
+ `tests/comm/test_ulysses_head_chunk.py` 16 passed / 4 skipped, and the
pcie_ipc module tests that import the JIT helpers from
`flashinfer.comm.pcie_ipc_ar` still pass; `pre-commit run --from-ref
<main> --to-ref HEAD` is green.

Re-run at the rebased head `0f60c090` on the same node (all eight GPUs
for the communicator suite, GPUs 0-3 for the rest while another tenant
held GPU 6): `tests/comm/test_ulysses_communicator.py` 133 passed / 22
skipped / 0 failed, `tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed,
`tests/comm/test_cake_ulysses_a2a.py` + `tests/comm/test_ulysses_a2a.py`
+ `tests/comm/test_ulysses_head_chunk.py` 15 passed / 5 skipped,
pcie_ipc module tests pass; `pre-commit run --from-ref <main> --to-ref
HEAD` is green.

On a single-GPU host the same files give 186 passed / 90 skipped
(host-mock and single-rank halves), `pre-commit run --all-files` is
green, `scripts/check_pr_document.py` reports no findings, and
`scripts/check_pr_api_diff.py` reports only the `max_elems` ->
`max_bytes` constructor advisory explained above.

Re-run at the rebased head `f0ff368c1` (2026-10-15) on the same node
with all eight GPUs, from a fresh checkout in a new directory (torch
2.14.1+cu130, nvcc 13.0, cold JIT cache):
`tests/comm/test_ulysses_communicator.py` 133 passed / 22 skipped / 0
failed, `tests/comm/test_ulysses_topology.py` 112 passed,
`tests/trace/test_ulysses_trace.py` +
`tests/trace/test_template_registry.py` 10 passed,
`tests/comm/test_cake_ulysses_a2a.py` + `tests/comm/test_ulysses_a2a.py`
+ `tests/comm/test_ulysses_head_chunk.py` 16 passed / 4 skipped, and the
six `tests/comm/test_pcie_ipc_*.py` files that import
`flashinfer.comm.pcie_ipc_ar` 73 passed / 12 skipped / 1 xfailed. The 22
skips are the forced-NVLink family; every PCIe case ran, including the
eight-rank all-RDMA and 4+4 hybrid routes. On the single-GPU host the
four core files give 187 passed / 90 skipped / 0 failed, `pre-commit run
--all-files` is green, `scripts/check_pr_document.py` reports no
findings and `scripts/check_pr_api_diff.py` only the `max_bytes`
advisory. `git range-diff` against the previous head `313ac10b8` shows
the 29 commits replay with context-only changes.

## Reviewer Notes

- The riskiest pieces are the epoch-barrier protocol
(`include/flashinfer/comm/ulysses_pcie.cuh`) and the mlx5 transport
(`csrc/ulysses_pcie_transport.cuh`): QP bring-up, interleaved-MKey
geometry (pitch ≤ 65535), completion accounting with per-exchange
immediate tags, and the abort/quiesce teardown ledger.
- `_validate`'s byte-budget capacity math after the per-call dtype
change (`flashinfer/comm/ulysses.py`) is worth a close read; the rescale
rationale is in the commit message of the dtype commit.
- The failure envelope is route-dependent on purpose, and the design
note's "Failure envelope" table is the contract to review against.
Python-side rejections are rank-uniform under the SPMD contract, so the
RDMA routes keep the communicator usable (the multi-rank tests rely on
refusing capture and a batch > 1 operand recoverably); the all-P2P route
poisons teardown because its barrier has no abort protocol at all.
Native failures on the RDMA routes publish the abort (now including the
per-call operand checks) and `close()` completes. Publishing the abort
from Python on every route is the first follow-up above.
- Things that are correct but will draw a question: (1) full-group CUDA
IPC is required even on the RDMA routes, because the barrier signals are
always IPC-mapped -- RDMA removes the P2P *data* requirement, not the
P2P requirement; (2) `batch == 1` on the RDMA routes comes from mode 0's
flat `rank * payload` remote offset, not from the NIC; (3) `input_ready`
is deliberately reused between `RunHybridBarrier` and `EnqueueCopies` so
the side streams inherit the post-barrier state; (4) the copy-engine
writes are ordered against the closing barrier's release store through
the `copy_done` event join, i.e. through `cudaMemcpyAsync` completion
semantics; (5) the per-call `dtype` is really per-`allocate_output`: the
registered buffer pins it and native re-checks every call against it.
- Placement: the backend uses the experimental machinery (warn-once,
gated `auto`, no AOT) but lives in `flashinfer/comm/` with tests in
`tests/comm/`, because it extends the stable `UlyssesCommunicator` class
and the routing shares its decision layer; there is no precedent for an
`@experimental_backend` in core. If the maintainers prefer the
Experimental Track layout, the code splits at the `_pcie_*` seam (the
methods touch only `rank`, `world_size`, `device`, `dtype`, `max_bytes`,
`backend`, `transport`, `decision`, `_state` and the `_gather` /
`_require_open` / `_validate` / `_output_geometry` helpers) and we will
move it.

---------

Co-authored-by: Alex Yang <aleyang@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants