Repository navigation
docs: improve Ulysses communicator and MoE EP docs - #4240
Merged
aleozlx merged 5 commits intoJul 31, 2026
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
kangbintNV
requested review from
Anerudhan,
Aneureka,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 29, 2026 13:43
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe changes add MoE EP environment guidance and expand communication API documentation. They document Ulysses communicator, topology, raw all-to-all, and ChangesDocumentation and API metadata updates
Estimated code review effort: 1 (Trivial) | ~5 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/api/comm.rst`:
- Around line 272-276: Move the autodoc directives for resolve_ulysses_backend,
decide_ulysses_backend, and probe_ulysses_rank_topology out of the “Raw Kernel
Entry Points (advanced)” section and place them before that heading or under a
dedicated topology-functions subsection, preserving all three functions in the
generated API navigation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 46749a59-79af-405c-bc40-c661ba303dbe
📥 Commits
Reviewing files that changed from the base of the PR and between 8b6af1a and 81611439389940b68d0c11589256cf97b70a0bb2.
📒 Files selected for processing (3)
CLAUDE.mddocs/api/comm.rstflashinfer/comm/ulysses.py
kangbintNV
changed the title
kangbintNV
force-pushed
the
kb/fix-doc-check-10-issues-20260729
branch
from
July 31, 2026 05:54
8161143 to
630f973
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/comm/mixed_comm.py`:
- Around line 938-945: Update the parameter documentation for comm_backend in
checkpoint_restore to state that Get_rank() must match local_rank and Get_size()
must match local_size. Clarify that the backend coordinates Unix-socket handle
exchange through bcast and barrier, rather than directly exchanging VMM
mappings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 6a814adf-0c92-454c-b773-ee8e5bcffacf
📥 Commits
Reviewing files that changed from the base of the PR and between 81611439389940b68d0c11589256cf97b70a0bb2 and 630f973.
📒 Files selected for processing (5)
CLAUDE.mddocs/api/comm.rstflashinfer/comm/mixed_comm.pyflashinfer/comm/ulysses.pyflashinfer/comm/ulysses_topology.py
🚧 Files skipped from review as they are similar to previous changes (3)
- CLAUDE.md
- docs/api/comm.rst
- flashinfer/comm/ulysses.py
aleozlx
reviewed
Jul 31, 2026
| UlyssesRankTopology | ||
| UlyssesBackendError | ||
|
|
||
| .. autofunction:: resolve_ulysses_backend |
aleozlx
pushed a commit
that referenced
this pull request
Aug 4, 2026
## 📌 Description Fixes the NEW documentation issues reported by the comm doc check. Changes: - document MoE EP environment variables in `CLAUDE.md` - keep Ulysses topology helper functions documented under the topology API section without marking them as `@flashinfer_api` - add missing parameter documentation for Ulysses communicator/raw A2A APIs - add missing parameter documentation for `MixedCommHandler.checkpoint_restore` ## 🔍 Related Issues N/A ## 🚀 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). - [x] 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.). Validation: - `git diff --check` - `python3 -m compileall flashinfer/comm/ulysses.py flashinfer/comm/ulysses_topology.py flashinfer/comm/mixed_comm.py` - AST check for the reported decorator/docstring/env-var items ## Reviewer Notes The Ulysses topology helper functions are intentionally documented but remain internal and are not marked with `@flashinfer_api`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added quick-reference documentation for three MoE EP environment variables. * Expanded Ulysses communication API documentation with parameters, tensor shapes, backend modes, device and dtype requirements, and return semantics. * Clarified topology probing, backend selection, and checkpoint restoration behavior. * Improved API reference generation for Ulysses backend utilities. <!-- end of auto-generated comment: release notes by coderabbit.ai --> (cherry picked from commit 668a1ba)
4 of 16 tasks
yyihuang
added a commit
that referenced
this pull request
Oct 12, 2026
…RMSNorm + RoPE overlapped with the Ulysses head exchange (SM90 / SM100 / SM103) (#6332) ## 📌 Description Part of #4254. Implements the overlap requested in #6304 (assumptions stated below where the issue thread has not confirmed a choice). Adds `flashinfer.comm.QkvNormRopeExchangeOverlap`: one prepared session that fuses the attention front end of a sequence-parallel (Ulysses) layer — the BF16 QKV projection, per-head RMSNorm of Q and K, interleaved RoPE from a per-token cos/sin table, and the head exchange — into a single call per layer. The Cake backend stores every destination's head band straight from the GEMM epilogue into that rank's symmetric receive slot over NVLink and publishes per-chunk readiness words, so the exchange overlaps the projection instead of following it; a one-CTA wait kernel on a consumer stream acquires the words and the caller's stream joins it. * Architectures: SM90 (WGMMA producer), SM100 and SM103 (tcgen05 producer; 2-CTA pair tile for `block_n = 256` and rows of at least 4096 tokens). Single-node NVLink groups of 2, 4 or 8 ranks, `head_dim = 128`, BF16 operands, fp32 cos/sin table. * Layouts: `"bshd"` returns `(q, k, v)` views of `[B, S_global, H_local, D]`; `"fused"` returns one `[B, S_global, H_local, 3 * D]` payload. Outputs are views of the current receive slot and stay valid until the slot is reused `num_slots` calls later; `out=` copies into caller-owned tensors. * Backend selection: `backend="auto"` goes through `resolve_ulysses_backend` (NVLink mesh probe) and the device capability; otherwise the NCCL path (cuBLAS projection, `flashinfer.norm.rmsnorm`, RoPE, `UlyssesCommunicator.scatter_qkv_head_chunk`). `backend="nvlink"` raises `UlyssesBackendError` when the mesh or the architecture is not supported; the NCCL path reports `graph_capturable == False`. * CUDA graphs: a prepared call allocates nothing and runs no host collective; captured calls publish a sentinel epoch and clear their readiness words with a stream memset after the wait kernel, so replays interleave with eager calls. A call with more rows than `max_tokens`, or a row count first seen inside a capture, raises instead of allocating. At most `num_slots` calls of one session can be captured into one graph without an eager call between. * `heads % world_size != 0` raises `ValueError` deterministically on every rank before any collective (same contract as `UlyssesCommunicator.scatter_heads`). ### Assumptions (open points in #6304) The issue thread has not yet confirmed the following choices; this PR implements the defaults below and labels them here so reviewers and the requester can object: * **Exchange algorithm**: Ulysses-style QKV all-to-all (each rank holds a token shard and the full QKV weight; every rank receives its head band of all tokens). Ring-style KV exchange is out of scope. * **Operator math**: the public FLUX definitions — fp32 RMSNorm with `eps = 1e-6` and a shared `[head_dim]` weight per Q / K, interleaved (pairwise) RoPE applied in fp32 from a per-token cos/sin table, in the order projection → QK-norm → RoPE; BF16 operands and outputs (no FP8 producer). * **Receive layout**: both the `(q, k, v)` `[B, S_global, H_local, D]` views and the fused `[B*S_global, 3*H_local*D]` payload are supported; the numbers below are measured with the `bshd` layout on every architecture. * **Tolerance**: `atol = rtol = 1e-2` against an fp32 step-by-step reference. Isolated single-ulp BF16 differences (a value landing on the other side of a BF16 rounding boundary after a different fp32 accumulation order, max_abs 0.03125) occur both in this kernel and in the existing per-operator path; the tests classify those elements explicitly instead of loosening the tolerance. ### Files | Path | Content | |---|---| | `flashinfer/comm/qkv_norm_rope_exchange_overlap.py` | public class, request validation, NCCL path | | `flashinfer/comm/cake_qkv_norm_rope_exchange_overlap.py` | Cake backend session: symmetric slots, barrier -> producer -> wait per call, capture protocol | | `flashinfer/comm/cake_qkv_exchange_plan.py` | tile plan and plan policy (host only) | | `flashinfer/jit/cake_qkv_norm_rope_exchange_overlap.py` | JIT loader: program and route tables, `gen_jit_spec` per program and architecture | | `csrc/cake_qkv_norm_rope_exchange_overlap/` | generated device and binding translation units (one `run` launcher per program; excluded from pre-commit like the other generated trees) | | `tests/comm/test_qkv_norm_rope_exchange_overlap.py` | host-only plan / validation / loader tests; 2-, 4-, 8-GPU correctness, capture/replay, refusals, NCCL path | | `benchmarks/comm/bench_qkv_norm_rope_exchange_overlap.py` | torchrun benchmark against the per-operator paths (CUPTI, cold L2) | | `docs/api/comm.rst`, `flashinfer/comm/__init__.py` | docs section and re-export | ## 🔍 Related Issues - #6304 — the feature request this PR implements; the choices the thread has not confirmed yet are labelled under "Assumptions" above. - Part of #4254. ## 🚀 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`. - [x] 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/). The hooks were run as `pre-commit run --files <the 31 files changed by this PR>` with the hook versions pinned in `.pre-commit-config.yaml` (unchanged by this PR); all 12 hooks pass, including clang-format, mypy, ruff check and ruff format. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Every path is checked against an fp32 step-by-step reference (projection, RMSNorm with the bf16 cast before the bf16 scale, interleaved RoPE, bf16 output) redistributed with `all_gather`, `atol = rtol = 1e-2`, on both layouts including a tail row count (`S_local = 200`), after CUDA-graph replays with fresh inputs, with two sessions in flight on one group, and on the NCCL path. The multi-GPU tests spawn 2, 4 and 8 ranks on one node and skip on fewer devices or an unsupported architecture. The host-only tests of `tests/comm/test_qkv_norm_rope_exchange_overlap.py` (`-k "not multi_gpu"`, 21 passed) ran on the pre-check node; the 2-, 4- and 8-GPU tests need a multi-GPU node and are left to CI. ### Test plan - [ ] `pytest tests/comm/test_qkv_norm_rope_exchange_overlap.py -q` on an 8-GPU SM100 node - [ ] same on an 8-GPU SM103 node - [ ] same on an 8-GPU SM90 node - [ ] `torchrun --standalone --nproc-per-node=8 benchmarks/comm/bench_qkv_norm_rope_exchange_overlap.py --rows representative --json bench_w8.json` per architecture - [x] `pre-commit run --files <new files>` (31 files; 12 / 12 hooks pass) ## 🔬 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 device and binding translation units under `csrc/cake_qkv_norm_rope_exchange_overlap/` are generated and attested by content hash; the host path (`flashinfer/comm/qkv_norm_rope_exchange_overlap.py`, `cake_qkv_norm_rope_exchange_overlap.py`, `cake_qkv_exchange_plan.py`), the JIT loader and the tests are hand-written and are the parts to read first. - The Speedup section separates the benchmark table from the export-protocol receipts; rows that failed the protocol's validity gates (power-capped clocks) are listed, not reported as measured. ## Speedup Measured with `benchmarks/comm/bench_qkv_norm_rope_exchange_overlap.py` (CUPTI activity timing, cold L2, per-rank maximum then median; `--rows representative`) against the fastest per-operator path on the same node. The per-operator baselines are: B1 three `scatter_heads` collectives, B2 one `scatter_qkv_head_chunk` over the whole head band, B3 the default head-chunk schedule, B4 `ulysses_a2a_in(fuse_qkv=True)` over NCCL. On SM90 the Ulysses communicator of B1-B3 runs on its NCCL backend (the NVLink Ulysses kernels target SM100 / SM103). ### Measured speedup (72 rows per architecture) Grid: `W ∈ {2, 4, 8}` ranks on one node × `hidden ∈ {3072, 4096, 6144, 8192}` (`head_dim = 128`) × `S_global ∈ {4096 … 131072}`, `B = 1`, BF16. Each row was timed on one 8-GPU node with CUPTI activity tracing (cold L2, per-rank maximum then median, five A/B/B/A rounds of ten samples, an A/A control arm within ±1 % or the row was re-measured) against the fastest of the four per-operator paths above measured in the same process. speedup = fastest baseline / fused. Every measured row is faster than its fastest baseline. | GPU | rows faster than the fastest baseline | geometric mean (A/A-valid rows) | W = 8 / 4 / 2 | minimum | |---|---|---|---|---| | B200 (SM100) | 72 / 72 | 1.831 | 1.824 / 1.849 / 1.820 | 1.261 | | B300 (SM103) | 72 / 72 | 1.906 | 1.911 / 1.891 / 1.916 | 1.241 | | H100 (SM90) | 71 / 71 | 1.332 | 1.339 / 1.348 / 1.310 | 1.135 | Notes: on B300 the row `W = 4, hidden 3072, S_global 8192` (2.26–2.31× in three measurements) exceeded the ±1 % A/A control in all three and is excluded from the geometric mean; on H100 (80 GB) the row `W = 2, hidden 8192, S_global 131072` does not fit in device memory with the reference and baseline buffers and is not measured. <details> <summary>B200 (SM100) — per-row table</summary> | W | hidden | S_global | fused ms | fastest baseline ms | speedup | A/A % | |---|---|---|---|---|---|---| | 8 | 3072 | 4096 | 0.0722 | 0.2389 | 3.307 | 0.11 | | 8 | 3072 | 8192 | 0.1014 | 0.2493 | 2.458 | 0.88 | | 8 | 3072 | 16384 | 0.1562 | 0.2821 | 1.807 | 0.09 | | 8 | 3072 | 32768 | 0.2713 | 0.4512 | 1.663 | 0.20 | | 8 | 3072 | 65536 | 0.5131 | 0.8815 | 1.718 | 0.10 | | 8 | 3072 | 131072 | 0.9575 | 1.7075 | 1.783 | 0.14 | | 8 | 4096 | 4096 | 0.0968 | 0.2440 | 2.522 | 0.71 | | 8 | 4096 | 8192 | 0.1193 | 0.2576 | 2.160 | 0.34 | | 8 | 4096 | 16384 | 0.1834 | 0.3626 | 1.977 | 0.67 | | 8 | 4096 | 32768 | 0.3086 | 0.6454 | 2.091 | 0.25 | | 8 | 4096 | 65536 | 0.5906 | 1.2448 | 2.108 | 0.05 | | 8 | 4096 | 131072 | 1.1808 | 2.4446 | 2.070 | 0.11 | | 8 | 6144 | 4096 | 0.1224 | 0.2503 | 2.045 | 0.62 | | 8 | 6144 | 8192 | 0.1902 | 0.3208 | 1.686 | 0.69 | | 8 | 6144 | 16384 | 0.3327 | 0.5789 | 1.740 | 0.32 | | 8 | 6144 | 32768 | 0.6438 | 1.1486 | 1.784 | 0.06 | | 8 | 6144 | 65536 | 1.3623 | 2.3209 | 1.704 | 0.17 | | 8 | 6144 | 131072 | 2.8806 | 4.4242 | 1.536 | 0.01 | | 8 | 8192 | 4096 | 0.1897 | 0.2766 | 1.458 | 0.26 | | 8 | 8192 | 8192 | 0.3278 | 0.4798 | 1.464 | 0.54 | | 8 | 8192 | 16384 | 0.5804 | 0.9063 | 1.562 | 0.22 | | 8 | 8192 | 32768 | 1.2319 | 1.7849 | 1.449 | 0.08 | | 8 | 8192 | 65536 | 2.5323 | 3.5145 | 1.388 | 0.06 | | 8 | 8192 | 131072 | 5.1799 | 6.9348 | 1.339 | 0.00 | | 4 | 3072 | 4096 | 0.0943 | 0.2456 | 2.605 | 0.25 | | 4 | 3072 | 8192 | 0.1316 | 0.2755 | 2.093 | 0.67 | | 4 | 3072 | 16384 | 0.2077 | 0.4497 | 2.165 | 0.31 | | 4 | 3072 | 32768 | 0.3794 | 0.8772 | 2.312 | 0.19 | | 4 | 3072 | 65536 | 0.6863 | 1.6978 | 2.474 | 0.06 | | 4 | 3072 | 131072 | 1.4040 | 3.3245 | 2.368 | 0.10 | | 4 | 4096 | 4096 | 0.1200 | 0.2506 | 2.088 | 0.23 | | 4 | 4096 | 8192 | 0.1768 | 0.3601 | 2.037 | 0.50 | | 4 | 4096 | 16384 | 0.3100 | 0.6342 | 2.046 | 0.34 | | 4 | 4096 | 32768 | 0.5834 | 1.2666 | 2.171 | 0.13 | | 4 | 4096 | 65536 | 1.1508 | 2.4901 | 2.164 | 0.01 | | 4 | 4096 | 131072 | 2.3583 | 4.9367 | 2.093 | 0.01 | | 4 | 6144 | 4096 | 0.1901 | 0.3145 | 1.654 | 0.87 | | 4 | 6144 | 8192 | 0.3349 | 0.5856 | 1.748 | 0.37 | | 4 | 6144 | 16384 | 0.6606 | 1.1643 | 1.763 | 0.13 | | 4 | 6144 | 32768 | 1.3538 | 2.3620 | 1.745 | 0.04 | | 4 | 6144 | 65536 | 2.8794 | 4.4967 | 1.562 | 0.14 | | 4 | 6144 | 131072 | 5.8146 | 8.9359 | 1.537 | 0.01 | | 4 | 8192 | 4096 | 0.3220 | 0.4808 | 1.493 | 0.05 | | 4 | 8192 | 8192 | 0.5756 | 0.9014 | 1.566 | 0.03 | | 4 | 8192 | 16384 | 1.2447 | 1.8077 | 1.452 | 0.03 | | 4 | 8192 | 32768 | 2.5719 | 3.5700 | 1.388 | 0.10 | | 4 | 8192 | 65536 | 5.2743 | 7.0276 | 1.332 | 0.01 | | 4 | 8192 | 131072 | 10.2687 | 14.3224 | 1.395 | 0.23 | | 2 | 3072 | 4096 | 0.1234 | 0.2666 | 2.160 | 0.92 | | 2 | 3072 | 8192 | 0.1972 | 0.4282 | 2.171 | 0.52 | | 2 | 3072 | 16384 | 0.3494 | 0.8643 | 2.474 | 0.04 | | 2 | 3072 | 32768 | 0.6568 | 1.6784 | 2.555 | 0.16 | | 2 | 3072 | 65536 | 1.3126 | 3.2883 | 2.505 | 0.20 | | 2 | 3072 | 131072 | 2.7374 | 6.5161 | 2.380 | 0.13 | | 2 | 4096 | 4096 | 0.1788 | 0.3531 | 1.975 | 0.69 | | 2 | 4096 | 8192 | 0.3104 | 0.6411 | 2.065 | 0.40 | | 2 | 4096 | 16384 | 0.5834 | 1.2492 | 2.141 | 0.11 | | 2 | 4096 | 32768 | 1.1547 | 2.4569 | 2.128 | 0.06 | | 2 | 4096 | 65536 | 2.3984 | 4.8868 | 2.038 | 0.10 | | 2 | 4096 | 131072 | 4.8963 | 9.6901 | 1.979 | 0.01 | | 2 | 6144 | 4096 | 0.3260 | 0.5658 | 1.735 | 0.40 | | 2 | 6144 | 8192 | 0.6637 | 1.1501 | 1.733 | 0.14 | | 2 | 6144 | 16384 | 1.3643 | 2.3287 | 1.707 | 0.01 | | 2 | 6144 | 32768 | 2.8530 | 4.4548 | 1.561 | 0.14 | | 2 | 6144 | 65536 | 5.7498 | 8.8343 | 1.536 | 0.26 | | 2 | 6144 | 131072 | 11.7871 | 17.7952 | 1.510 | 0.15 | | 2 | 8192 | 4096 | 0.5740 | 0.8991 | 1.566 | 0.21 | | 2 | 8192 | 8192 | 1.2396 | 1.7892 | 1.443 | 0.01 | | 2 | 8192 | 16384 | 2.5594 | 3.5306 | 1.379 | 0.08 | | 2 | 8192 | 32768 | 5.2429 | 6.9682 | 1.329 | 0.23 | | 2 | 8192 | 65536 | 10.5083 | 14.1865 | 1.350 | 0.21 | | 2 | 8192 | 131072 | 23.0629 | 29.0725 | 1.261 | 0.63 | </details> <details> <summary>B300 (SM103) — per-row table</summary> | W | hidden | S_global | fused ms | fastest baseline ms | speedup | A/A % | |---|---|---|---|---|---|---| | 8 | 3072 | 4096 | 0.0774 | 0.2697 | 3.487 | 0.42 | | 8 | 3072 | 8192 | 0.1027 | 0.2863 | 2.788 | 0.66 | | 8 | 3072 | 16384 | 0.1579 | 0.3131 | 1.982 | 0.88 | | 8 | 3072 | 32768 | 0.2759 | 0.4525 | 1.640 | 0.39 | | 8 | 3072 | 65536 | 0.5072 | 0.8430 | 1.662 | 0.14 | | 8 | 3072 | 131072 | 0.9269 | 1.6214 | 1.749 | 0.02 | | 8 | 4096 | 4096 | 0.0982 | 0.2790 | 2.842 | 0.36 | | 8 | 4096 | 8192 | 0.1223 | 0.2896 | 2.367 | 0.54 | | 8 | 4096 | 16384 | 0.1788 | 0.3498 | 1.956 | 0.49 | | 8 | 4096 | 32768 | 0.2904 | 0.6136 | 2.113 | 0.04 | | 8 | 4096 | 65536 | 0.5417 | 1.1756 | 2.170 | 0.12 | | 8 | 4096 | 131072 | 1.0635 | 2.3010 | 2.164 | 0.05 | | 8 | 6144 | 4096 | 0.1258 | 0.2846 | 2.263 | 0.61 | | 8 | 6144 | 8192 | 0.1809 | 0.3167 | 1.751 | 0.51 | | 8 | 6144 | 16384 | 0.3056 | 0.5489 | 1.797 | 0.12 | | 8 | 6144 | 32768 | 0.5970 | 1.0776 | 1.805 | 0.01 | | 8 | 6144 | 65536 | 1.2046 | 2.0970 | 1.741 | 0.06 | | 8 | 6144 | 131072 | 2.4691 | 4.1316 | 1.673 | 0.01 | | 8 | 8192 | 4096 | 0.1852 | 0.2919 | 1.576 | 0.01 | | 8 | 8192 | 8192 | 0.3020 | 0.4518 | 1.496 | 0.74 | | 8 | 8192 | 16384 | 0.5334 | 0.8476 | 1.589 | 0.19 | | 8 | 8192 | 32768 | 1.0841 | 1.6653 | 1.536 | 0.10 | | 8 | 8192 | 65536 | 2.1971 | 3.2764 | 1.491 | 0.07 | | 8 | 8192 | 131072 | 4.4165 | 6.4274 | 1.455 | 0.03 | | 4 | 3072 | 4096 | 0.0945 | 0.2622 | 2.774 | 0.69 | | 4 | 3072 | 8192 | 0.1346 | 0.3073 | 2.283 (A/A > 1 %) | 4.33 | | 4 | 3072 | 16384 | 0.2104 | 0.4479 | 2.128 | 0.17 | | 4 | 3072 | 32768 | 0.3899 | 0.8362 | 2.145 | 0.33 | | 4 | 3072 | 65536 | 0.6896 | 1.6060 | 2.329 | 0.11 | | 4 | 3072 | 131072 | 1.3541 | 3.1379 | 2.317 | 0.05 | | 4 | 4096 | 4096 | 0.1207 | 0.2752 | 2.280 | 0.66 | | 4 | 4096 | 8192 | 0.1735 | 0.3476 | 2.003 | 0.09 | | 4 | 4096 | 16384 | 0.2905 | 0.6134 | 2.111 | 0.01 | | 4 | 4096 | 32768 | 0.5332 | 1.1986 | 2.248 | 0.06 | | 4 | 4096 | 65536 | 1.0389 | 2.3463 | 2.258 | 0.02 | | 4 | 4096 | 131072 | 2.0955 | 4.6322 | 2.211 | 0.01 | | 4 | 6144 | 4096 | 0.1834 | 0.3102 | 1.691 | 0.39 | | 4 | 6144 | 8192 | 0.3131 | 0.5571 | 1.779 | 0.09 | | 4 | 6144 | 16384 | 0.6136 | 1.0952 | 1.785 | 0.33 | | 4 | 6144 | 32768 | 1.2197 | 2.1381 | 1.753 | 0.01 | | 4 | 6144 | 65536 | 2.4356 | 4.1972 | 1.723 | 0.06 | | 4 | 6144 | 131072 | 4.9588 | 8.3107 | 1.676 | 0.03 | | 4 | 8192 | 4096 | 0.3039 | 0.4521 | 1.488 | 0.15 | | 4 | 8192 | 8192 | 0.5342 | 0.8521 | 1.595 | 0.06 | | 4 | 8192 | 16384 | 1.1044 | 1.6901 | 1.530 | 0.09 | | 4 | 8192 | 32768 | 2.1950 | 3.3179 | 1.512 | 0.08 | | 4 | 8192 | 65536 | 4.4062 | 6.5251 | 1.481 | 0.10 | | 4 | 8192 | 131072 | 9.4225 | 13.3648 | 1.418 | 0.28 | | 2 | 3072 | 4096 | 0.1199 | 0.2954 | 2.463 | 0.11 | | 2 | 3072 | 8192 | 0.1977 | 0.4372 | 2.211 | 0.15 | | 2 | 3072 | 16384 | 0.3375 | 0.8262 | 2.448 | 0.06 | | 2 | 3072 | 32768 | 0.5991 | 1.5968 | 2.665 | 0.04 | | 2 | 3072 | 65536 | 1.1957 | 3.1269 | 2.615 | 0.01 | | 2 | 3072 | 131072 | 2.3627 | 6.1780 | 2.615 | 0.10 | | 2 | 4096 | 4096 | 0.1690 | 0.3380 | 1.999 | 0.06 | | 2 | 4096 | 8192 | 0.2898 | 0.6124 | 2.113 | 0.01 | | 2 | 4096 | 16384 | 0.5334 | 1.1898 | 2.230 | 0.10 | | 2 | 4096 | 32768 | 1.0432 | 2.3309 | 2.234 | 0.09 | | 2 | 4096 | 65536 | 2.0717 | 4.6073 | 2.224 | 0.01 | | 2 | 4096 | 131072 | 4.2197 | 9.1462 | 2.167 | 0.05 | | 2 | 6144 | 4096 | 0.2977 | 0.5423 | 1.822 | 0.04 | | 2 | 6144 | 8192 | 0.6118 | 1.0884 | 1.779 | 0.04 | | 2 | 6144 | 16384 | 1.1907 | 2.1221 | 1.782 | 0.02 | | 2 | 6144 | 32768 | 2.3919 | 4.1900 | 1.752 | 0.04 | | 2 | 6144 | 65536 | 4.9604 | 8.2729 | 1.668 | 0.34 | | 2 | 6144 | 131072 | 11.4356 | 16.7756 | 1.467 | 0.23 | | 2 | 8192 | 4096 | 0.5286 | 0.8535 | 1.615 | 0.11 | | 2 | 8192 | 8192 | 1.0660 | 1.6769 | 1.573 | 0.01 | | 2 | 8192 | 16384 | 2.1406 | 3.3157 | 1.549 | 0.22 | | 2 | 8192 | 32768 | 4.5065 | 6.5023 | 1.443 | 0.24 | | 2 | 8192 | 65536 | 9.5439 | 13.3589 | 1.400 | 0.16 | | 2 | 8192 | 131072 | 22.1867 | 27.5295 | 1.241 | 0.03 | </details> <details> <summary>H100 (SM90) — per-row table</summary> | W | hidden | S_global | fused ms | fastest baseline ms | speedup | A/A % | |---|---|---|---|---|---|---| | 8 | 3072 | 4096 | 0.1070 | 0.1708 | 1.595 | 0.54 | | 8 | 3072 | 8192 | 0.1587 | 0.2288 | 1.442 | 0.13 | | 8 | 3072 | 16384 | 0.2722 | 0.3818 | 1.403 | 0.23 | | 8 | 3072 | 32768 | 0.5272 | 0.7191 | 1.364 | 0.16 | | 8 | 3072 | 65536 | 1.0614 | 1.4336 | 1.351 | 0.57 | | 8 | 3072 | 131072 | 2.1269 | 2.8805 | 1.354 | 0.34 | | 8 | 4096 | 4096 | 0.1444 | 0.1875 | 1.298 | 0.33 | | 8 | 4096 | 8192 | 0.2046 | 0.3137 | 1.533 | 0.45 | | 8 | 4096 | 16384 | 0.3804 | 0.5538 | 1.456 | 0.29 | | 8 | 4096 | 32768 | 0.7378 | 1.1132 | 1.509 | 0.87 | | 8 | 4096 | 65536 | 1.4582 | 2.1800 | 1.495 | 0.03 | | 8 | 4096 | 131072 | 2.9126 | 4.3244 | 1.485 | 0.13 | | 8 | 6144 | 4096 | 0.2603 | 0.3047 | 1.171 | 0.50 | | 8 | 6144 | 8192 | 0.4252 | 0.5419 | 1.274 | 0.27 | | 8 | 6144 | 16384 | 0.7434 | 1.0295 | 1.385 | 0.33 | | 8 | 6144 | 32768 | 1.5361 | 2.0466 | 1.332 | 0.46 | | 8 | 6144 | 65536 | 3.2696 | 4.0454 | 1.237 | 0.11 | | 8 | 6144 | 131072 | 6.4700 | 8.0205 | 1.240 | 0.01 | | 8 | 8192 | 4096 | 0.3395 | 0.4407 | 1.298 | 0.41 | | 8 | 8192 | 8192 | 0.6529 | 0.8201 | 1.256 | 0.36 | | 8 | 8192 | 16384 | 1.3057 | 1.6042 | 1.229 | 0.16 | | 8 | 8192 | 32768 | 2.6329 | 3.2563 | 1.237 | 0.08 | | 8 | 8192 | 65536 | 5.4181 | 6.4155 | 1.184 | 0.25 | | 8 | 8192 | 131072 | 11.4591 | 13.0004 | 1.135 | 0.23 | | 4 | 3072 | 4096 | 0.1570 | 0.2114 | 1.347 | 0.18 | | 4 | 3072 | 8192 | 0.2545 | 0.3688 | 1.449 | 0.02 | | 4 | 3072 | 16384 | 0.4623 | 0.7117 | 1.540 | 0.05 | | 4 | 3072 | 32768 | 0.8927 | 1.3835 | 1.550 | 0.02 | | 4 | 3072 | 65536 | 1.7488 | 2.7224 | 1.557 | 0.08 | | 4 | 3072 | 131072 | 3.4939 | 5.3476 | 1.531 | 0.11 | | 4 | 4096 | 4096 | 0.1987 | 0.2856 | 1.437 | 0.02 | | 4 | 4096 | 8192 | 0.3767 | 0.5468 | 1.452 | 0.09 | | 4 | 4096 | 16384 | 0.7412 | 1.0631 | 1.434 | 0.06 | | 4 | 4096 | 32768 | 1.4287 | 2.0803 | 1.456 | 0.01 | | 4 | 4096 | 65536 | 2.8435 | 4.1326 | 1.453 | 0.02 | | 4 | 4096 | 131072 | 5.7662 | 8.1828 | 1.419 | 0.27 | | 4 | 6144 | 4096 | 0.4104 | 0.5240 | 1.277 | 0.08 | | 4 | 6144 | 8192 | 0.7647 | 1.0093 | 1.320 | 0.07 | | 4 | 6144 | 16384 | 1.4989 | 1.9741 | 1.317 | 0.11 | | 4 | 6144 | 32768 | 3.0037 | 3.9102 | 1.302 | 0.05 | | 4 | 6144 | 65536 | 6.1564 | 7.8220 | 1.271 | 0.11 | | 4 | 6144 | 131072 | 13.0916 | 15.5928 | 1.191 | 0.02 | | 4 | 8192 | 4096 | 0.6562 | 0.8133 | 1.239 | 0.12 | | 4 | 8192 | 8192 | 1.2991 | 1.5901 | 1.224 | 0.00 | | 4 | 8192 | 16384 | 2.5754 | 3.1704 | 1.231 | 0.00 | | 4 | 8192 | 32768 | 5.2416 | 6.2735 | 1.197 | 0.35 | | 4 | 8192 | 65536 | 11.1591 | 12.7948 | 1.147 | 0.10 | | 4 | 8192 | 131072 | 22.6522 | 26.1679 | 1.155 | 0.07 | | 2 | 3072 | 4096 | 0.2538 | 0.3503 | 1.380 | 0.19 | | 2 | 3072 | 8192 | 0.4533 | 0.6655 | 1.468 | 0.04 | | 2 | 3072 | 16384 | 0.8781 | 1.2965 | 1.476 | 0.06 | | 2 | 3072 | 32768 | 1.7103 | 2.5706 | 1.503 | 0.00 | | 2 | 3072 | 65536 | 3.4275 | 4.9401 | 1.441 | 0.03 | | 2 | 3072 | 131072 | 7.0964 | 9.8401 | 1.387 | 0.17 | | 2 | 4096 | 4096 | 0.3736 | 0.5182 | 1.387 | 0.19 | | 2 | 4096 | 8192 | 0.7116 | 1.0057 | 1.413 | 0.06 | | 2 | 4096 | 16384 | 1.3972 | 1.9768 | 1.415 | 0.00 | | 2 | 4096 | 32768 | 2.7507 | 3.9208 | 1.425 | 0.02 | | 2 | 4096 | 65536 | 5.5682 | 7.6357 | 1.371 | 0.06 | | 2 | 4096 | 131072 | 11.8074 | 15.2643 | 1.293 | 0.08 | | 2 | 6144 | 4096 | 0.7606 | 0.9606 | 1.263 | 0.06 | | 2 | 6144 | 8192 | 1.4823 | 1.8973 | 1.280 | 0.02 | | 2 | 6144 | 16384 | 2.9507 | 3.7534 | 1.272 | 0.09 | | 2 | 6144 | 32768 | 5.9493 | 7.3609 | 1.237 | 0.13 | | 2 | 6144 | 65536 | 12.6436 | 14.8480 | 1.174 | 0.12 | | 2 | 6144 | 131072 | 26.1080 | 30.5084 | 1.169 | 0.46 | | 2 | 8192 | 4096 | 1.2689 | 1.5342 | 1.209 | 0.09 | | 2 | 8192 | 8192 | 2.5373 | 3.0496 | 1.202 | 0.12 | | 2 | 8192 | 16384 | 5.1135 | 6.0219 | 1.178 | 0.22 | | 2 | 8192 | 32768 | 10.6399 | 12.2423 | 1.151 | 0.22 | | 2 | 8192 | 65536 | 21.6914 | 25.0241 | 1.154 | 0.07 | | 2 | 8192 | 131072 | — | — | not measured (device memory) | — | </details> ### Export-protocol validation (per-row receipts, three architectures) Independently of the table above, every row was re-measured by the export protocol that produced the generated sources in this PR: same CUPTI method (cold L2, per-rank maximum then median, counterbalanced A/B/B/A groups), the generated programs loaded through this PR's JIT path, the four per-operator baselines measured in the same process, plus the validity gates per row (source/export agreement >= 0.95, directional disagreement <= 0.12, endpoint drift <= 0.15, loaded SM clock >= 50 % of the maximum boost clock). A row counts only when all gates pass. **Protocol change (2026-10-11): loaded-SM-clock floor 0.60 -> 0.50.** The floor used to be 0.60. At the largest `W = 2 / 4, S_global = 131072` rows the B200 (1000 W) and B300 (1100 W) board power limits hold the loaded SM clock at 0.48 - 0.60 of the maximum on every node tried (five nodes per SKU, three rounds), so those rows were rejected unconditionally. The floor is now 0.50; the directional, drift and source gates are unchanged. Every previously accepted row was measured under the stricter 0.60 floor and remains valid as measured; under the relaxed floor one further row was accepted (`corr_w2_max` on B300, a correctness row). The rows that still fail are listed below. | GPU | perf rows with a valid receipt | all valid rows faster than the fastest baseline | geometric mean (W = 8 / 4 / 2) | minimum | rows without a valid receipt | |---|---|---|---|---|---| | B200 (SM100) | 69 / 72 | yes | 1.686 (1.535 / 1.716 / 1.832) | 1.170 | 3 | | B300 (SM103) | 69 / 72 | yes | 1.700 (1.511 / 1.737 / 1.889) | 1.109 | 3 | | H100 (SM90) | 70 / 72 | yes | 1.306 (1.296 / 1.319 / 1.301) | 1.093 | 2 | Correctness / negative rows of the export protocol (13 per GPU, including the largest-shape correctness row `corr_w2_max`): H100 13 / 13, B300 13 / 13, B200 12 / 13 (`corr_w2_max` reads 0.496 and 0.492 of the maximum clock on two nodes, below the 0.50 floor; its correctness is covered by the harness correctness table above). Rows without a valid receipt (disclosed, not counted; the raw ratio of an invalid attempt is not a measurement): | GPU | W | hidden | S_global | why the receipt is invalid | attempts (nodes) | raw ratio of the last attempt | |---|---|---|---|---|---|---| | B200 | 2 | 8192 | 131072 | rounds 1 - 2: loaded SM clock 847 / 1170 / 1147 MHz of 1965 under the 1000 W limit; round 3 at the 0.50 floor: B1 endpoint drift 0.30 / 0.22 and B2 directional disagreement > 0.12 | 4 (3 nodes) | 1.37 - 1.69 | | B200 | 4 | 8192 | 131072 | rounds 1 - 2: loaded SM clock below the 60 % floor (960 / 1177 MHz); round 3: B1 directional disagreement > 0.12 | 3 (3 nodes) | 1.55 - 1.89 | | B200 | 2 | 6144 | 131072 | directional disagreement of the B1 / B3 baselines > 0.12 (candidate gates pass) | 2 (2 nodes) | 1.54 - 1.90 | | B300 | 2 | 8192 | 131072 | rounds 1 - 2: loaded SM clock 1117 / 1140 / 1170 MHz of 2032 under the 1100 W limit; round 3: B2 / B3 directional disagreement > 0.12 | 4 (3 nodes) | 1.44 - 1.84 | | B300 | 4 | 8192 | 131072 | rounds 1 - 2: loaded SM clock 1147 / 1215 MHz of 2032; round 3: B1 directional disagreement, B2 directional disagreement + endpoint drift | 4 (3 nodes) | 1.47 - 1.90 | | B300 | 2 | 6144 | 131072 | directional disagreement of the B2 / B3 baselines > 0.12 | 3 (2 nodes) | 1.70 - 2.16 | | H100 | 8 | 8192 | 131072 | directional disagreement of the B4 (NCCL) baseline 0.15 - 0.18 > 0.12 in all five attempts (candidate and B1 - B3 gates pass); the raw B2 ratio of the attempts that reached the arms is 0.97 - 0.99 | 5 (5 nodes) | 0.97 - 1.15 | | H100 | 2 | 8192 | 131072 | does not fit in 80 GB with the reference and the four baseline arms (not measurable) | - | - | The largest shapes run power-capped on the B200 / B300 nodes available for this work (clocks at 48 - 60 % of the maximum boost); once the clock floor no longer rejects them, the order effect between the baseline arms (each tens of milliseconds, with power-state transitions between arms) exceeds the 0.12 directional gate on every node tried. The per-operator table above measured those rows on other nodes of the same type, where the row `W = 2, hidden 8192, S_global 131072` reached 23 ms vs 28.9 ms. The export protocol is not relaxed further for them; they are listed here rather than reported as measured. ### Baselines and their source PRs | Baseline | Path | Source PRs | |---|---|---| | B1 ordinary | cuBLAS projection → `flashinfer.norm.rmsnorm` → `apply_rope_with_cos_sin_cache` → `UlyssesCommunicator.scatter_heads` × 3 | #3820, #4240 (communicator), #4480, #4202 (rmsnorm), #5952, #5956, #6188 (RoPE) | | B2 whole-QKV head chunk | one `scatter_qkv_head_chunk` collective over the whole head band | #5027 | | B3 head-chunk schedule | `scatter_qkv_head_chunk` per head band on three streams, the band schedule of `benchmarks/comm/bench_ulysses_head_chunk.py` (default schedule per row) | #5027, #5263, #5953 | | B4 NCCL all-to-all | `parallel_wrapper.ulysses_a2a_in(fuse_qkv=True)` | #2630 | Baseline implementation files (all at the merge-base): `flashinfer/comm/ulysses.py` (B1 `scatter_heads`, B2 and B3 `scatter_qkv_head_chunk`, and `UlyssesCommunicator`), `benchmarks/comm/bench_ulysses_head_chunk.py` (B3 band schedule), `flashinfer/norm/` (B1 `rmsnorm`), `flashinfer/rope.py` (B1 `apply_rope_with_cos_sin_cache`), `flashinfer/parallel_attention/parallel_wrapper.py` (B4 `ulysses_a2a_in`). This branch is based on `34bb66e4` (its merge-base with `main`); `git log 34bb66e..main -- <those files>` is empty as of `main` = `65c38d2a`, so the baselines measured here are the implementations currently in `main`. Source PR merge commits: #3820 `8a94642d8`, #4240 `668a1ba1c`, #4480 `6d5ef77f8`, #4202 `519205914`, #5952 `a973cfc92`, #5956 `335c4eb4e`, #6188 `d71cc4b38`, #5027 `15e83b7bb`, #5263 `8c54f3446`, #5953 `7e73a7470`, #2630 `ce02358d2`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a prepared sequence-parallel API combining QKV projection, Q/K normalization, rotary embeddings, and head exchange. * Automatically selects the optimized NVLink backend when supported, with an NCCL fallback for other configurations. * Supports CUDA graph capture on the optimized backend and reusable output slots. * **Documentation** * Added API guidance on supported configurations, backend selection, outputs, and capture behavior. * **Benchmarks** * Added a distributed benchmark comparing the fused operation with separate operator paths, including output checks and optional JSON results. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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
Fixes the NEW documentation issues reported by the comm doc check.
Changes:
CLAUDE.md@flashinfer_apiMixedCommHandler.checkpoint_restore🔍 Related Issues
N/A
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Validation:
git diff --checkpython3 -m compileall flashinfer/comm/ulysses.py flashinfer/comm/ulysses_topology.py flashinfer/comm/mixed_comm.pyReviewer Notes
The Ulysses topology helper functions are intentionally documented but remain internal and are not marked with
@flashinfer_api.Summary by CodeRabbit