Repository navigation
Conversation
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the 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 |
Contributor
Summary of ChangesHello @Amir-19, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request streamlines the memory allocation process for NVLink all-reduce and all-reduce fusion operations by adopting PyTorch's symmetric memory API. This change unifies the underlying memory management for both single-node and multi-node scenarios, simplifying the codebase and leveraging a robust, built-in PyTorch mechanism for distributed memory. The refactoring ensures proper resource handling and sets the stage for more consistent and efficient communication primitives. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request unifies memory allocation for single-node and multi-node NVLink all-reduce operations by leveraging torch.distributed._symmetric_memory. The changes introduce a new module torch_symmetric_memory.py for symmetric buffer allocation and refactor existing all-reduce workspace creation and destruction functions to use this new mechanism. This is a positive step towards simplifying memory management and better integrating with PyTorch's distributed capabilities. However, there are a few areas related to type annotations, handling of optional parameters, and consistent use of dtype that could be improved for robustness and clarity.
| lamport_buffer_size = tp_size * LamportTokenNumThreshold * tp_size * hidden_dim * 2 | ||
|
|
||
| device = torch.device(f"cuda:{torch.cuda.current_device()}") | ||
| group_name = group.group_name |
Contributor
There was a problem hiding this comment.
The group_name is assigned group.group_name directly. However, the group parameter is Optional[ProcessGroup]. If group is None, this line will raise an AttributeError. It should handle the None case gracefully, similar to how it's done in trtllm_create_ipc_workspace_for_all_reduce_fusion (line 598).
Suggested change
| group_name = group.group_name | |
| group_name = group.group_name if group is not None else torch.distributed.group.WORLD.group_name |
|
|
||
| def _alloc_symm_buffer_bytes( | ||
| size_bytes: int, tp_size: int, dtype: torch.dtype, device: torch.device, group_name: str | ||
| ) -> tuple[list[int], torch.Tensor]: |
Contributor
There was a problem hiding this comment.
The return type annotation for _alloc_symm_buffer_bytes is tuple[list[int], torch.Tensor], but the function actually returns three values: ptrs, tensor, and handle. The handle is of type torch.distributed._symmetric_memory.SymmetricMemoryHandle. Please update the type hint to accurately reflect the return signature.
Suggested change
| ) -> tuple[list[int], torch.Tensor]: | |
| ) -> tuple[list[int], torch.Tensor, symm_mem.SymmetricMemoryHandle]: |
| mapping.tp_rank, | ||
| torch.device("cuda", mapping.local_rank), | ||
| comm_backend, | ||
| torch.float32, #TODO(asamani): should this always be f32 or dtype? |
Contributor
There was a problem hiding this comment.
The dtype for _alloc_symm_buffer_bytes is hardcoded to torch.float32. The MNNVLAllReduceFusionWorkspace constructor takes a dtype parameter. For consistency and to support different data types, this should use the dtype passed to the constructor.
Suggested change
| torch.float32, #TODO(asamani): should this always be f32 or dtype? | |
| dtype, |
|
|
||
| # We use FP32 for sentinel value regardless of the real dtype | ||
| self.mcast_buffer_handle.lamport_initialize(mapping.tp_rank, torch.float32) | ||
| self.lamport_initialize(mapping.tp_rank, torch.float32, allocated_size) |
Contributor
There was a problem hiding this comment.
The dtype for lamport_initialize is hardcoded to torch.float32. This should ideally use the dtype passed to the MNNVLAllReduceFusionWorkspace constructor to ensure consistency with the allocated buffer's data type.
Suggested change
| self.lamport_initialize(mapping.tp_rank, torch.float32, allocated_size) | |
| self.lamport_initialize(mapping.tp_rank, dtype, allocated_size) |
Contributor
yyihuang
added a commit
that referenced
this pull request
Oct 23, 2026
…_norm_combine, cake_ulysses_backend): one source per clone group, one binding, one prelude header (#7054) > Developed by the **CAKE team**, the kernels in this PR **outperform** the relevant **SOTA baselines at the time of merge** on the **vast majority of performance benchmark shapes**, matching them otherwise. ## Summary This PR is the round-2 deslop of five Cake comm units for SM100 / SM103, tracked in #6674. The units are `cake_moe_alltoall`, `cake_trtllm_moe_allreduce_union`, `cake_moe_finalize_allreduce_fusion`, `cake_fused_norm_combine` and `cake_ulysses_backend`. Commit 1 compiles both all-to-all targets from one byte-identical source. Commit 2 builds each of three union pipe clone groups from one kept source through compile definitions. Commit 3 drops dead route typing from the union loader and unread tables from the finalize loader. Commit 4 shares one host binding across the three norm-combine variants. Commit 5 moves the shared device prelude of the 24 Ulysses kernel TUs into one header. Every former kernel symbol stays defined for its module. Device code, module names, routes, launch parameters, message texts and packaging do not change. Behaviour changes: - **Union loader surface.** `exported_arches()` is removed. Its only caller was one test assertion, which is now written against `ROUTES`. `Route.kernel` is always one kernel name, because the per-rank tuple form had no route. The `Kernel` record gains two keyword fields: `source` (the kept source stem; `None` means the record's own name) and `defines` (its compile definitions). Both have defaults, so every positional literal stays valid. `kernel_source(name)` and `kernel_defines(name)` honour them. - **Finalize loader surface.** The unread tables `WORLD_SIZES`, `DTYPES` and `OUTPUT_PROFILES` and the helper `kernel_names()` are removed from the module and its `__all__`. `HIDDEN_DIM` stays because the benchmark reads it. (The public `SUPPORTED_COMPUTE_CAPABILITIES` of the Ulysses package also stays.) - **Norm-combine.** `VARIANTS[*]` gains a `shim` field, which holds the host-shim namespace each former binding spelled. When a host-side check in the fused norm-combine binding fails, the traceback file / line now names the shared binding (its path and line numbers) instead of the deleted per-variant file. The error kind, the message text and every check predicate are unchanged. The kernel symbol inside the launch-failure message is spelled through `CAKE_STRINGIZE` and comes out as the former literal. - **One rebuild on an existing JIT workspace.** Module names stay the same. Every module whose source text or source path changed rebuilds once: - the 2 all-to-all modules (kernel source path); - per architecture, the 7 union modules that compile the 3 kept pipe sources (3 kept + 4 re-pointed); - per architecture, the 3 norm-combine modules (shared binding); - per architecture, the finalize module (one comment line of its binding); - per architecture, the 24 Ulysses libraries (prelude include). A fresh workspace is not affected. The 48 `closure_sha256` values of the Ulysses `build.json` are unchanged, so the Ulysses library names do not change. No in-tree code reproduces the exporter's closure recipe, the manifest check requires a 64-hex value, and the spec-name suffix stays stable. - **No device code changes.** For every kernel symbol, the gates compile both sides to identical SASS from the loaders' exact compile definitions (see Validation). The kept MNNVL source is byte-identical to both former files. ## Changes Base: upstream main `0602a5b99`. The branch was cut from `f2f5add99`, and the counts "at the base" below were taken there. Head: `b2fae49ab2d264bed7ed458fc5ba5c344dc40f7f` (5 commits). 45 files, +394 / -9,771 (1 added, 7 deleted, 35 modified, 2 renamed). - **Units.** This PR touches five Cake comm units, all served only on SM100 (B200, compute capability 10.0) and SM103 (GB300 / B300, compute capability 10.3). Each one is JIT-compiled lazily per module and architecture, either through its loader under `flashinfer/jit/` or through the Ulysses package's own `cake_build` loader. - `cake_moe_alltoall`, the MoE all-to-all over MNNVL (loader `flashinfer/jit/comm.py`, generated CUDA under `csrc/generated/moe_alltoall/`). It has two JIT modules, `cake_alltoall_sm100a` / `cake_alltoall_sm103a`. Each is compiled from the host sources plus one generated kernel source. - `cake_trtllm_moe_allreduce_union`, the MoE all-reduce "union" of 27 generated kernel programs. Its `KERNELS` / `ROUTES` tables are in `flashinfer/jit/cake_trtllm_moe_allreduce_union.py`, its kernel sources are under `csrc/cake_trtllm_moe_allreduce_union/kernels/`, and it has one shared launcher TU. It builds 27 modules per architecture. - `cake_moe_finalize_allreduce_fusion`, the finalize + all-reduce fusion (loader `flashinfer/jit/cake_moe_finalize_comm.py`, host binding `csrc/cake_moe_finalize_allreduce_fusion/cake_moe_finalize_allreduce_fusion_binding.cu`). It builds one module per architecture. - `cake_fused_norm_combine`, the fused norm + combine. It has 3 `VARIANTS` in `flashinfer/jit/cake_fused_norm_combine.py`, with kernel TUs and host bindings under `csrc/cake_fused_norm_combine/`. It builds 3 modules per architecture. - `cake_ulysses_backend`, the Ulysses all-to-all package `flashinfer/cake_ulysses_backend/`. It has 24 kernel TUs and 24 bindings under its `csrc/`, plus a `build.json` manifest with 48 library entries (24 × 2 architectures), and it is loaded through `cake_build`. It builds 24 modules per architecture. - **Sharing form.** This PR uses the fold-through-definitions form that is already merged for the Cake MLA families. - A kept generated kernel source spells the tokens that differed between its former clones through compile definitions. Their guarded defaults equal the file's own former values. As a result, the record that keeps the file builds with its exact former command line and the former token stream. - The loader record of each folded member names the kept source and carries the member's own values as compile definitions. Only that source receives them. - A kept host binding works the same way for the launcher-side differences: kernel symbol, shim namespace, block size, dynamic shared memory and cooperative launch. - Fold arms are literal: `#if NAME == <kept value>`, the former kept line, `#else`, the former member line, `#endif`. They never sit inside `__launch_bounds__` or `#pragma unroll`, and they are never keyed on `__CUDA_ARCH__`. - **Unchanged.** Every former kernel symbol stays defined for its module. Every module name, `KERNELS` key, `ROUTES` entry, `VARIANTS` key, launch parameter, message text and accepted input is unchanged. Packaging is unchanged. - **Generated files.** Every edited generated file carries this one-line statement exactly once, directly after the licence block: `// Generated by Cake for FlashInfer. Exporter output, hand-edited here pending exporter support.` That covers 29 files: the 3 kept union kernel sources, the shared norm-combine binding (where the statement replaces the "do not edit" banner), the 24 Ulysses kernel TUs and the new Ulysses prelude header. The deletions and the byte-identical rename in commit 1 edit no generated text. `build.json` is exporter data. The finalize binding, the union launcher, the Python loaders and the tests are hand-written files and carry no statement. - **Census corrections to the issue text.** The issue was written before two upstream changes. At the base, the issue's 59 kernel sources (the five units plus the two DCP all-to-all kernels) are 68,367 lines. The issue's figure of 36,269 lines counts the former `csrc/generated/ulysses/` set of 24 files / 4,982 lines. #6864 replaced that set with the 24 package TUs (37,080 lines) under `flashinfer/cake_ulysses_backend/csrc/`. #6621 renamed the MNNVL sources to `csrc/generated/moe_alltoall/cake_moe_alltoall_sm100.cu` and `cake_moe_alltoall_sm103.cu`, which the loader table in `flashinfer/jit/comm.py` selects. ### Commit 1 (77a3fe3) - cake_moe_alltoall: compile both Blackwell targets from one source (4 files, +6 / -6,545) - **Duplicate.** The two generated kernel sources, `csrc/generated/moe_alltoall/cake_moe_alltoall_sm100.cu` and `cake_moe_alltoall_sm103.cu`, are byte-identical: 6,533 lines, one sha256, one git blob. The exporter emitted the same text twice. That text selects its one architecture-specific statement with a `__CUDA_ARCH__` guard. - **One file.** The SM100 copy is renamed to `cake_moe_alltoall_kernels.cu` without changing a byte, and the SM103 copy is deleted. Both targets of the loader table in `flashinfer/jit/comm.py` now point at the one file. - **Unchanged:** the module names `cake_alltoall_sm100a` / `cake_alltoall_sm103a`, the per-target architecture flags, the other sources and the AOT registration. Each target compiles its own module from the same bytes, as before. - **Test.** The all-to-all structure test's four file-name pins follow the rename. Its 33-kernel, signature and definition checks run unchanged on the kept file. - **Disclosure:** none, because no generated text is edited. ### Commit 2 (6cc1808) - cake_trtllm_moe_allreduce_union: build the pipe clones from one source (9 files, +146 / -2,018) - **Clones.** Three groups of generated pipe kernels differ only in three literals each: the launch-bounds minimum-blocks value, the pipe depth and the unroll factor. - **Kept sources.** One source is kept per group: `ws2_bf16_pipe2_u4_b5`, `ws2_f16_pipe2_u4_b5` and `ws8_bf16_pipe1_u4_b5`. For example, `ws2_bf16_pipe2_u4_b5` means world size 2, BF16, pipe depth 2, unroll 4 and min blocks 5. - The differing literals go through `CAKE_UNION_MIN_BLOCKS`, `CAKE_UNION_PIPE_DEPTH` and `CAKE_UNION_UNROLL`, whose guarded defaults equal the kept member's values. Every differing site has literal arms, and no macro sits inside `__launch_bounds__` or `#pragma unroll`. - The kernel symbol line goes through `CAKE_UNION_KERNEL`. The launcher already receives that definition for every module. - The three kept sources carry the disclosure line. - **Deleted.** The four clone sources are deleted: `ws2_bf16_pipe1_u4_b5`, `ws2_bf16_pipe2_u4`, `ws2_f16_pipe1_u4_b5` and `ws8_bf16_pipe1`. Together they had 2,006 lines at the base (454 + 454 + 448 + 650). - **Loader.** The four folded `KERNELS` records name the kept source and carry their values as compile definitions. `KERNELS` keys, `ROUTES`, spec names and launch parameters are unchanged. - **Test.** The kernel-symbol pin accepts the `CAKE_UNION_KERNEL(` spelling when the module's definitions name the symbol. - **Not folded:** the ws4 pair, which is benchmark-gated, and the bf16 / f16 twins, which are not clones. ### Commit 3 (2af24a6) - cake comm loaders: drop dead route typing and unused finalize tables (4 files, +15 / -45) - **Union loader.** `Route.kernel` becomes `str`. The tuple form, and its per-rank branch in the kernel-name helper, had no route. The unused `Union` import is removed. `exported_arches()` and its `__all__` entry are removed, and the test now derives the architecture set from `ROUTES`. Two docstring sentences that described the tuple form are reworded to match the code. - **Finalize loader.** The unread `WORLD_SIZES` / `DTYPES` / `OUTPUT_PROFILES` tables, `kernel_names()` and their `__all__` entries are removed. `HIDDEN_DIM` stays because the benchmark reads it. - **Not removed:** the Ulysses package's `SUPPORTED_COMPUTE_CAPABILITIES`, a public constant. - **Finalize binding.** One comment line now names the current kernel source file instead of a deleted path. The comment keeps its line count, so no `__LINE__` of its `TVM_FFI_CHECK` sites moves. - **Consequence:** no kernel source, module name or flag changes. ### Commit 4 (8106de7) - cake_fused_norm_combine: share one host launcher across the variants (4 files, +71 / -251) - **Clones.** The three per-variant host bindings have the same text except for five things: the kernel symbol, the host-shim namespace, the block size, the dynamic shared-memory size, and whether the launch is cooperative. - **Kept binding.** One binding is kept: `csrc/cake_fused_norm_combine/cake_fused_norm_combine_binding.cu`. Its text comes from the former `84f760a6` binding. Git's rename detection pairs the new file with the former `490c897e` binding instead, so the rename shown in the diff does not identify where the text came from. - The five differences go through `CAKE_FUSED_NORM_COMBINE_KERNEL`, `CAKE_FUSED_NORM_COMBINE_SHIM`, `CAKE_FUSED_NORM_COMBINE_BLOCK`, `CAKE_FUSED_NORM_COMBINE_DYN_SMEM` and `CAKE_FUSED_NORM_COMBINE_COOPERATIVE`. These definitions are strict: a missing definition is a compile error. - The symbol inside the launch-failure message is spelled through `CAKE_STRINGIZE`. - The disclosure line replaces the former "do not edit" banner. The two other bindings are deleted. - **Loader.** `VARIANTS[*]["sources"][1]` names the shared binding, and each variant gains its `shim` field. The loader passes the five definitions to the binding alone through `extra_cuda_cflags_by_source`. Every kernel TU's command line is therefore byte-identical. - **Unchanged:** `capacity_receipt` and `COMPILE_FLAGS`, which are exporter-filled data that the tests read as an oracle, and the exports (`run` and `max_active_blocks_per_sm`). ### Commit 5 (b2fae49) - cake_ulysses_backend: move the device prelude into one header (26 files, +156 / -912) - **Prelude.** All 24 generated kernel TUs of the Ulysses package open with the same 36-line device prelude. - **Header.** The new `flashinfer/cake_ulysses_backend/csrc/cake_ulysses_a2a_prelude.cuh` holds the licence block, that exact prelude text and the disclosure line. In every TU, the prelude block is replaced by `#include "cake_ulysses_a2a_prelude.cuh"` plus the disclosure line. The package `csrc` is already on the include path. - **`build.json`.** Exactly the 48 kernel-TU `sha256` values (24 TUs × 2 architecture entries) are regenerated as the digests of the new file bytes. The `sources` lists, the 48 binding digests and all 48 `closure_sha256` values are unchanged, so every library name is unchanged. - **Not here:** the Ulysses dtype fold (24 to 8 TUs, 24 to 3 bindings). ## Validation These rows are compile-only: the SASS gate, the MNNVL rows, the preprocessed-text gate, the host gate and the launcher inertness rows. They run on one CPU node with one CUDA toolkit for both architectures, and the SASS gate's result gives the toolkit version. Every row compares the base tree with the branch head. Both sides of a row are compiled from one staged path in one environment with the loader's exact flag list. The GPU rows run on three nodes: a B200 node (SM100, 8 GPUs), a GB300 node (SM103, 4 GPUs) and a B300 node (SM103, 8 GPUs). Both trees use fresh JIT workspaces. - SASS gate, 68 rows (34 per architecture: the 7 union pipe kernel sources touched by the fold (14 rows: the 3 kept sources with their guarded defaults and the 4 folded members under their compile definitions, each against its former source at the base), the 3 norm-combine kernel TUs (6), the 24 Ulysses kernel TUs (48); cubin + SASS + kernel set + resource usage): all 68 rows identical (68 distinct compiles per side; the 0 coincident kept-route pairs compiled once and reported under both names; per unit U2 14 / U3 6 / U5a 48; per architecture sm_100a 34 / sm_103a 34): both sides of every row compiled from one staged path in one environment with the loader's exact flag list, the normalised SASS identical for every former kernel symbol (no symbol mask), exactly the expected kernel function set in every cubin, resource usage identical; nvcc Build cuda_13.4.r13.4/compiler.38855100_0 (CUDA 13.4) - **Pass criteria.** A row passes when all of these hold: - the SASS is identical for every former kernel symbol, with no symbol mapping; - every cubin has exactly the expected kernel function set; - resource usage is identical. - **Scope.** The other 20 union kernel sources are byte-identical to the base, and they appear in the module build-and-load row. The SASS identity of the all-to-all kernel source is covered by the 2 MNNVL rows below, which are counted separately from the 68. - MNNVL single source, 2 rows (per architecture the base file against the kept file): all 2 rows identical: the kept generated source carries the digest of both former per-architecture files, and for each architecture the base file and the kept file compiled from one staged path give identical normalised SASS with exactly the 33 kernel functions; the device-pass preprocessed text identical after dropping line markers (the renamed basename is the only raw difference) - The kept source must hash to the digest of both former files. The SASS and kernel set must be identical per architecture. - Preprocessed-text gate, 62 rows (the device pass of every surviving kernel source at head against its former source at the base): all 62 rows identical: the device-pass preprocessed text of every surviving kernel source at head against its former source at the base (taken from the compile of the SASS rows) identical after dropping line markers and blank lines, with no whitespace, literal or symbol allowance - The texts must be identical after dropping line markers and blank lines. No other difference is allowed. - Host gate, 8 rows (the shared norm-combine binding under each variant's definitions against that variant's former binding (6 rows: identical after line markers / blank lines, the one adjacent-literal join at the launch-failure message and the check macros' file / line arguments with exactly the expected count per launcher); the finalize binding at head against the base (2 rows: exact after line markers / blank lines); compiled host objects): all 8 rows identical: the shared norm-combine launcher at head under each variant's compile definitions against that variant's former launcher at the base (6 rows: exact after line markers / blank lines, the one adjacent-literal join site at the launch-failure message (two joins per launcher) and the check macros' file / line arguments with exactly the expected count per launcher) and the finalize binding at head against the base (2 rows: exact after line markers / blank lines); the compiled host objects export the same symbols and reference the same undefined kernel symbol on both sides; no device code on either side - The compiled host objects must have identical export sets and reference the same undefined kernel symbol. Neither side may contain device code. - Launcher inertness, 8 rows (the union launcher TU compiled under each folded member's module flag list, its new definitions appended, against the kept record's list): all 8 rows identical: the union launcher compiled under each re-pointed member's module flag list (its new compile definitions appended) against the kept record's list, preprocessed host text identical after line markers / blank lines, export set identical - The preprocessed host text must be identical after dropping line markers / blank lines, and the export set must be identical. Together these show that the definitions are inert for the launcher. - Launch-bounds rows, 2 informational rows (the (224) vs (224, 5) pairs of the union family; nothing is dropped in this PR): 2 informational rows (nothing dropped in this PR): `lb__ws2_bf16_pipe2_u4__vs__pipe2_u4_b5__sm_100a`: masked SASS differs; `lb__ws8_bf16_pipe1__vs__pipe1_u4_b5__sm_103a`: masked SASS differs - The result states whether the masked SASS is identical. If it is, the owner may drop `ws2_bf16_pipe2_u4_b5` in a follow-up. - Module build and load through the real loaders (per architecture: 1 all-to-all + 27 union + 3 norm-combine + 1 finalize + 24 Ulysses modules, 56 per tree per architecture, 112 per tree; both trees, fresh JIT workspaces, on B200 and on the two SM103 nodes): 56 modules per tree for sm_100a / 56 modules per tree for sm_103a built and loaded through the real loaders in fresh JIT workspaces on both trees on B200 (SM100), GB300 (SM103), B300 (SM103): spec names, export sets and device kernel sets identical to the base-tree modules, every build cold, the re-pointed union members rebuilt from the kept source, the finalize loader's removed names raising on this branch only, the Ulysses kernel objects' dependency files naming the new header - Spec names, export sets and device kernel sets must be identical to those of the base-tree modules. - Every build must be cold, and the folded union members must be rebuilt from the kept source. - The removed names of the finalize loader must raise on this branch only. - The dependency files of the Ulysses kernel objects must name the new header. - Compile-line check (every gate row's flag list against the loader's generated build-file compile line of the corresponding object): every gate row's flag list equals the loader's generated build-file compile line of the corresponding object (continuation lines joined, common flags substituted) on B200 (SM100), GB300 (SM103), B300 (SM103) (35 / 35 / 35 rows) - The gate must compile exactly what the loader compiles. - Loader tables re-read on both trees (import with CUDA hidden; the union `KERNELS` / `ROUTES` / source and definitions per name / route selection over a token grid; the norm-combine `VARIANTS` with the added `shim` field and the variant selection; the all-to-all loader spec capture with the unknown-target rejection): the 3 resolved-table dumps (union KERNELS / ROUTES / kernel source and definitions per name / route selection over the token grid and the specialization rules; norm-combine VARIANTS with the added shim field and the variant selection sweep; the all-to-all loader spec capture with the unknown-target rejection) identical on both trees after masking the checkout roots, imported with CUDA hidden - After the checkout roots are masked, the tables must be identical apart from the declared deltas: the folded records' source / definitions and the `shim` field. - Tests on B200, GB300 and B300 (both the base and this branch, fresh JIT workspaces) and with CUDA hidden: B200 (SM100), 8 GPUs: mnnvl 203 passed / 0 skipped, union_dist 165 passed / 1 skipped, norm_combine 23 passed / 0 skipped, ulysses 18 passed / 0 skipped, finalize 11 passed / 0 skipped identical on both trees; GB300 (SM103), 4 GPUs: mnnvl 203 passed / 0 skipped, union_dist 163 passed / 3 skipped, norm_combine 22 passed / 1 skipped, ulysses 14 passed / 4 skipped, finalize 9 passed / 2 skipped identical on both trees; B300 (SM103), 8 GPUs: mnnvl 203 passed / 0 skipped, union_dist 165 passed / 1 skipped, norm_combine 23 passed / 0 skipped, ulysses 18 passed / 0 skipped, finalize 11 passed / 0 skipped identical on both trees (fresh JIT workspaces per tree; 15 route / group runs over 5 groups: mnnvl, union_dist, norm_combine, ulysses, finalize); per id the outcome and the skip reason are identical between the base and this branch, the id sets identical, 0 failed / 0 errors; skip classes: 4 skipped identically on both trees because fewer GPUs on the node than the cell's world size; 3 skipped identically on both trees because other: 'CUDA graph capture needs a GPU'; 1 skipped identically on both trees because other: 'requires eight SM100 or eight SM103 GPUs on one node'; MNNVL cells passed on B200 (SM100) (42 device cells passed), GB300 (SM103) (42 device cells passed), B300 (SM103) (42 device cells passed) - **The five groups:** - `tests/comm/test_cake_moe_allreduce_distributed.py`: world sizes 2 / 4 / 8 × fp16 / bf16. The ws 8 cells for SM103 run on the 8-GPU B300 node. - `tests/comm/test_cake_fused_norm_combine.py`: the 8-peer cell runs on the 8-GPU nodes, and the CPU items run everywhere. - `tests/comm/test_cake_ulysses_a2a.py`: ws 2 / 4 / 6 / 8 × kernel variants. The registry tests run everywhere. - `tests/comm/test_cake_moe_finalize_allreduce.py`: ws 2 / 4 / 8. - the MNNVL group: `tests/comm/test_cake_moe_alltoall.py` and `tests/comm/test_cake_moe_alltoall_fused_adapter.py` (two-rank MPI cells and structure tests), and `tests/comm/test_cake_moe_alltoall_multi_gpu.py` (four-rank MPI cells). - **Both trees.** - For each test id, the outcome and the skip reason must be identical on both trees. - The id sets must be identical, with 0 failed. - Cells whose world size exceeds a node's GPU count must skip identically on both trees. The result reports them. - **CUDA hidden.** These CPU companions must pass with CUDA hidden on both trees: - `tests/comm/test_cake_moe_allreduce_union.py`; - `tests/comm/test_cake_moe_allreduce_api.py`; - `tests/comm/test_cake_moe_finalize_dispatch.py`; - `tests/comm/test_cake_moe_alltoall_recv_view_cache.py`; - the structure and registry tests. - CPU checks: all 20 rows passed on a CPU node: pre-commit --all-files on both full checkouts (16 hooks, 15 passed / 0 failed / 1 skipped, identical per hook, empty patches); the five edits re-applied to a snapshot of the base reproduce every commit tree and the branch byte for byte (47 paths; the kept all-to-all source staged under its ignored path as one index entry), with every unit's text proofs equal to their recorded results and the proofs re-run at the head; the gate inventory regenerated from the tree equals the staged one; collection identical on both trees for the nine CI test files; the host-only items pass with CUDA hidden on both trees with identical outcomes; no deleted file name or removed loader name referenced anywhere; the kept all-to-all source hashes to both former files' digest and is the one entry under its directory with pairwise-distinct object names for both targets; the Ulysses manifest check passes and the 48 closure digests are unchanged with exactly the kernel-source digests regenerated; every edited generated file carries the disclosure statement once; commit message checks - pre-commit runs on both full checkouts, where it must give identical per-hook results and an empty patch. It also runs on the touched files. - Re-applying the five edits to a snapshot of the base must reproduce every commit's tree and the branch byte for byte. The kept all-to-all source lives under an ignored path and is staged as one index entry. Each unit's text proofs must equal their recorded results, and the proofs are re-run at the head. - The gate inventory regenerated from the tree must equal the staged one. - Collection must be identical on both trees for the nine CI test files. The host-only items must be identical with CUDA hidden. - No deleted file name or removed loader name may be referenced anywhere. - The kept all-to-all source must hash to the digest of both former files. It must be the only entry under its directory, and the object names for the two targets must be distinct. - The Ulysses manifest check must pass. The 48 closure digests must be unchanged, and exactly the kernel-source digests must be regenerated. - Every edited generated file must carry the disclosure statement once. - Commit message checks. - Type check: the repository's mypy hook passed on this branch and on the base, with identical findings - The repository's mypy hook runs on this branch and on the base, and both runs must give identical findings. **Not run.** By design, this PR's validation does not include the following items. Each item gives its proof of record. - **Benchmarks.** No benchmark of any family runs. This PR is mechanical and keeps the SASS of every kernel symbol identical, so there is nothing new to measure. The ws4 pair, the 18 exact-expert specializations and the two launch-bounds pairs are left as they are. They are benchmark decisions for the owner, and the launch-bounds rows above are informational. - **H100 / SM90.** No unit of this PR serves SM90, because every loader maps only compute capability 10.0 and 10.3. Nothing in this PR executes there. - **Two-node runs.** Every harness is single-host. The SM103 ws 8, 8-peer and ws 6 cells run on the 8-GPU B300 node instead. - **Cross-tree tensor comparison.** No comm test saves its outputs. The identity criterion is therefore identical outcomes per test id, plus the gate rows: identical SASS, preprocessed text, host text and export / kernel sets. ## Not in this PR The numbered exporter requests below are listed in the issue comment. - **The ws4 pipe pair and the 18 exact-expert union specializations:** benchmark-gated, so this is the owner's decision. - **Dropping `ws2_bf16_pipe2_u4_b5` / `ws8_bf16_pipe1_u4_b5` for their (224) twins:** this PR adds only the gate rows. Dropping them is a follow-up if the owner wants it. - **The Ulysses dtype fold** (24 to 8 kernel TUs, and 24 to 3 bindings with a world-size definition) **and the `MODULES` leaf dedup:** exporter request 6, or a hand follow-up. - **`SUPPORTED_COMPUTE_CAPABILITIES` of the Ulysses package:** a public constant, so it is kept. - **The DCP all-to-all pair fold** (`csrc/generated/dcp_alltoall/`) **and the DCP / finalize prelude headers:** the pair differs in renamed temporaries, not in compile definitions. These are exporter requests 5 / 4. - **One cross-family `cake_prelude.cuh`:** the per-family preludes differ in three details. This is exporter request 4. - **`capacity_receipt` / `COMPILE_FLAGS` of the norm-combine loader:** these are exporter-filled data that the tests read as an oracle. - **Host-path items** (`_check_cake_moe_allreduce_arch`, the loader re-checks, the launcher tensor checks, the finalize caches, `v_scale_factor`, the SM-count helper): these are behaviour / exporter / owner changes (#6672). - **The union benchmark and CI enrollment changes:** these are coverage additions. - **Every other comm family** (the #6675 files, `cake_fp4_allreduce`, `cake_gemm_allreduce`, `cake_fp8_rowwise_gemm_allreduce`, `cake_dsv4`, `cake_hopper_msa`, `cake_msa_icp_indexer`, the non-Cake comm code): these belong to other issues or are untouched. - **The standing-rule files** (`.pre-commit-config.yaml`, `scripts/check_cake_build_manifests.py`, `tests/test_helpers/cake_package.py`, any `cake_build.py`): untouched by rule. - **Concurrent upstream work:** no commit after the base touches these families. Some open PRs touch neighbouring comm files: - #4876 adds a function to `flashinfer/jit/comm.py` next to the loader table that commit 1 re-points; - #2389 / #4634 edit `flashinfer/comm/trtllm_ar.py`, which this PR does not touch; - #5317 touches `flashinfer/comm/ulysses_experimental.py`, a file that is not in this tree and that this PR does not touch. Part of #6674 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Refactor** * Consolidated several GPU kernel variants to use shared sources and configurable build settings. * Standardized shared setup across GPU communication kernels while keeping their implementations unchanged. * Updated kernel selection and routing to use a single kernel name per route. * **Chores** * Refreshed generated-kernel metadata to match the updated sources. * Updated tests to reflect the revised kernel selection and shared-source configuration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude <noreply@anthropic.com>
This branch has not been deployed
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
🔍 Related Issues
🚀 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.).Reviewer Notes