mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
0043f43309
## Summary Add a public, runtime-wide `liger_cute_kernels.configure(...)` API for shared NVSHMEM initialization and preparation of **multiple TP and EP groups**. This supplies the upstream state-management prerequisite discussed in https://github.com/verl-project/verl/pull/7945 without making runtime ownership specific to the output head. **Draft: single-node B300 and H200 runtime coverage is recorded below. Multi-host execution and downstream VeRL integration remain pending.** Rebased onto upstream main at `6ad077c`; the already merged CUDA-extras change is not included again. ```python from liger_kernel.ops.configure import FusedLinearCrossEntropyConfig, MoEConfig, configure native_configured = configure( process_groups={"tp": pg_collection.tp, "ep": pg_collection.ep}, bootstrap_group=world_group, device=device, flsce=FusedLinearCrossEntropyConfig( max_tokens=8192, hidden_size=4096, local_vocab_size=32000, group="tp", ), moe=MoEConfig( max_tokens=8192, hidden_size=4096, num_experts=128, top_k=8, num_hosts=ep_num_hosts, gpus_per_host=ep_gpus_per_host, max_inflight=8, # Reserve for all simultaneously retained MoE forwards. group="ep", ), ) ``` ## Configuration and group ownership - One canonical runtime implementation lives in LCK. `liger_kernel.ops.configure` is a thin optional entry point with dependency-free configuration dataclasses. Importing it and constructing configuration data do not require LCK. When LCK is absent, `configure` warns and returns `False`; successful setup returns `True`, and errors from an installed LCK propagate. - Bootstrap defaults to torch WORLD, not the first operator's group. Named groups accept Megatron's `pg_collection` without a Megatron dependency. TP and EP may use different uniform strided partitions of the same bootstrap domain, including a non-WORLD bootstrap domain. - Both sections can be configured together, added later, or repeated to prepare additional partitions and group sizes within their original process-wide capacities. Additional names are additive; an existing name retains its ordered membership. Equal/smaller requests reuse reservations rather than clearing another consumer's state. - Setup collectively validates capacities, topology, membership, cache agreement, CUDA hardware, and runtime ownership. Partial cache hits still participate in required team splits and symmetric reservations. Equivalent ordered memberships reuse a context. - The Python execution path caches `ProcessGroup -> team_handle`; execution never creates a team. Native FLSCE and MoE contexts select the prepared group, and queued kernels/captured graphs retain their own pointers and topology. - Both SM90 and SM100 forward/backward paths use the selected context. MoE visits peers in team-local round-robin order, `(rank + offset) % team_size`, local first. It has no mutable device-global schedule tables or schedule-specific host-topology fields in its kernel parameters/iterators. Non-communicating kernels need no group workspace; Triton vocab-parallel CE already uses caller-owned process groups and per-call tensors. ## Memory and lifecycle - Keep the existing buffer pool; there is no TP- or EP-specific pool wrapper. Fixed-size symmetric payload scratch is shared at the process-wide maximum. TP signaling/mapping state is context-specific. MoE device workspace is context-specific so another group's growth cannot invalidate captured pointers. - Symmetric slots follow bootstrap-wide allocation order, not local team-handle values. Cache-hit contexts retain their original allocations and epochs. Reservations remain resident through model offload and inference-engine sleep/wake. - MoE retained forward tensors stay in the symmetric stack, not the buffer pool. `MoEConfig.max_inflight` defaults to 1 and reserves a bound **per EP context**; count all layers and microbatches awaiting backward. Separate contexts have independent stacks; equivalent groups share their limit. Pop in LIFO order within a context. Overflow is rejected before launching or partially mutating stacks. - Low-level MoE metadata lookup/pop select the forward's team; the autograd wrapper forwards it automatically. Destroying a MoE team with live forward intermediates is rejected. Team destruction, pool clear, and finalization invalidate the relevant contexts and Python aliases. - Setup reserves symmetric storage before independent subgroup execution. Legacy lazy allocation now synchronizes preceding device collectives before a fresh WORLD heap allocation, fixing the stacked-forward deadlock found by the full suite. - Shared-scratch entry/exit synchronization prevents a later group from reusing data while an earlier group still consumes it. ## Additional fixes found during review and full-suite execution - Implement SM100 backward's node-local direct-peer path using the existing FP32 direct-peer all-reduce, including context-local epochs and full-tile scatter. The previous failure with `NVSHMEM_DISABLE_NVLS=1` is addressed rather than excluded. - Size Python MoE tile-ID storage for the minimum supported tile size. - Fix the native test build's missing SM90 contract include and exclude the standalone cuBLAS benchmark from the core library source glob. - Replace the obsolete skipped MoE roundtrip placeholder with a real single-PE eager/CUDA-graph forward/backward test. - Remove obsolete CUDA schedule symbols, host permutation builders/uploads, flat schedule-setup exports, and non-RDC module table plumbing. Preserve positive-topology validation and use a 64-bit product to reject overflow. Non-RDC setup initializes the CUDA context explicitly rather than relying on a table upload's side effect. Binding tests now bootstrap NVSHMEM before configuration, matching the native API contract. - Correct the MLP1 numerical reference to use independent cuBLAS BF16-input/FP32-output GEMMs followed by CPU activation arithmetic. The old sequential CPU GEMM disagreed near cancellation, including with FP64 accumulation; an independent cuBLAS probe reproduced the kernel's worst outputs exactly. **Original random inputs, shapes, 1% mean tolerance, and 5% maximum tolerance are unchanged.** Both fused and activation variants now have maximum relative error below 0.4% on B300. Matching native extensions must be rebuilt for the changed host/TVM binding interfaces. Dependency floors and the CUDA 12.9 release workflow are unchanged. ## MoE schedule performance and spilling The topology-arithmetic implementation was rejected after controlled B300 comparisons showed forward latency increases of 1.02-3.03% versus symbolic tables. The final implementation uses simple round-robin and removes the obsolete table infrastructure. Compared against symbolic-table controls on separate eight-GPU B300 and H200 nodes: EP2/EP8, four shapes, three counterbalanced process trials, 15 samples/trial, 20 native calls per captured batch. CUDA-event timings take the maximum across ranks; results below are medians of trial medians. The first trial checks outputs and gradients against Torch references outside timing. Clocks were not locked. H200 comparisons use **non-RDC for both variants**, with `LIGER_CUTE_ENABLE_SM90_NONRDC_MOE=ON`, `LIGER_CUTE_SM90_NONRDC_ALL_CONFIGS=ON`, and `LIGER_MOE_SM90_NONRDC=1`. All 66 local/IB main-kernel specializations are present in each matching cubin. Default-RDC Hopper measurements are not used to evaluate this change. | Single-node latency change vs. tables (negative is faster) | B300 | H200 non-RDC | | --- | --- | --- | | Forward | -2.30% to -0.39% | -1.47% to +0.29% | | Backward | -0.60% to +0.12% | -0.61% to -0.03% | | Full forward + backward | -0.76% to -0.19% | -0.57% to +0.08% | H200 is not uniformly faster: the largest EP2 forward shape has a +0.29% median increase (paired trials +0.29% to +1.44%), accepted by the author as within the performance-noise tolerance for this cleanup. One EP2 full-pass median increases +0.08%, with paired trials crossing zero. All EP8 H200 medians improve. B300's largest EP8 backward case has a small consistent +0.05% median increase; every B300 forward and full-pass paired trial improves. Compiler-reported spill changes for the final kernels: | Path | Spill-store / spill-load bytes vs. tables | | --- | --- | | B300 forward, 8 configs | 24/24 -> 20/20 | | B300 backward, 3 configs | Unchanged at 172/292 | | H200 non-RDC local forward, 10 configs | Unchanged at zero | | H200 non-RDC local backward, 23 configs | Unchanged | | H200 non-RDC IB-capable forward, 10 configs | -8 / -24 bytes | | H200 non-RDC IB-capable backward, 23 configs | -4 / -8 bytes | Spill counts are static ptxas metrics, not runtime memory traffic. `REG:168` is the kernel-level allocation, not the post-`setmaxnreg` MMA-warp limit: Hopper dynamically uses producer/consumer counts of 64/216 forward and 88/208 backward. Single-node timings exercise local transport; lower IB-capable spill counts do **not** establish multi-host HCA performance. ## Execution contract and remaining limits Use one CUDA device per process. Configure serially at a coordinated setup boundary, with every bootstrap rank participating, outside graph capture. Pipeline stages cannot invoke WORLD-collective setup only from post-process chunks; choose an appropriate common bootstrap domain. **Serialize all native launches and graph replays on each process, including across different TP/EP groups**, using one CUDA stream or explicit synchronization. Participating ranks must issue collectives in consistent order. Warm up shapes/configurations before capture. Finish dependent launches, backward passes, and graphs before destroying teams or clearing pools. Live capacity growth remains unsupported; reserve process-wide maxima up front. Legacy low-level MoE setup retains WORLD-collective lazy allocation; use public `configure` for independent EP subgroup execution. LCK rejects an already initialized visible unmanaged NVSHMEM runtime. DeepEP V1 RDMA ownership sharing remains unsupported in either initialization order. Detection does not guarantee interoperability with hidden runtimes or prevent another library from initializing after LCK. Skipping optional setup does not change backend selection or enable an explicitly selected native-only operator. ## Validation All implementation and review work was performed directly, without subagents. Environment: supplied eight-GPU B300 and H200 nodes, CUDA compiler **12.9.86**, Torch **2.9.1.6+cu129**. B300 uses Triton **3.5.1**, system NVSHMEM **3.6.0**, and an isolated system-site-packages venv. This does not validate the release-pinned NVSHMEM 3.6.5 runtime. H200 runs use non-RDC MoE, node-local NVSHMEM transport (`NVSHMEM_DISABLE_NCCL=1`, `NVSHMEM_REMOTE_TRANSPORT=none`, `NVSHMEM_SYMMETRIC_SIZE=6G`), and disable unrelated system pytest plugin autoload. | Final-source run | Result | | --- | --- | | SM100f native core and complete C++ test build | Passed | | SM90a non-RDC all-config cubin, matching native core, and C++ test build | Passed | | B300 complete LCK Python suite and Liger `test/cute` | 180 passed, no skips | | B300 complete C++ CTest suite | 66 passed, 26 hardware/opt-in benchmark skips, 0 failures | | Liger FLSCE frontend contracts | 21 passed | | B300 corrected NVSHMEM-bootstrap binding fixture | 12 passed | | H200 non-RDC bindings plus forward graph targeted retry | 13 passed | | H200 non-RDC schedule and all-dispatch native tests | 4 passed | | H200 non-RDC MoE binding, graph, multi-EP, and frontend suite | 27 passed, 1 architecture-gated frontend skip | | H200 complete C++ CTest suite, non-RDC MoE enabled | 55 passed, 30 hardware/opt-in benchmark skips, 0 failures | | Changed-file Ruff check/format and `git diff --check` | Passed | The complete suite commands are: ```bash PYTHONPATH=liger_cute_kernels:src uv run --no-project --python .venv/bin/python \ python -m pytest --confcutdir=test/transformers -v -o log_cli=false \ --tb=short --disable-warnings --durations=10 \ liger_cute_kernels/test test/cute ctest --test-dir build/sm100f --output-on-failure --timeout 180 -j 1 ``` Multi-group coverage includes 4/6/8 GPUs, overlapping groups, different group sizes, unequal prior launch counts, partial cache hits, split bootstrap domains, retained forwards across context switches, smaller workloads under reserved maxima, overflow rejection, graph replay, device-workspace growth in another EP group, team destruction/recreation, and pool-clear/reconfiguration. Existing mixed-operator coverage exercises distinct MoE EP and FLSCE TP groups. The first H200 Python run exposed missing CUDA-context initialization after table-upload removal and a binding fixture that configured before NVSHMEM initialization. Both were corrected without weakening assertions. Another case stalled in NVSHMEM bootstrap before entering kernels; rerunning with the explicit node-local transport settings above resolved it. These initial failures are not counted as successful runs. Multi-host communication, release-pinned NVSHMEM 3.6.5, DeepEP interoperability, and downstream VeRL changes are not claimed complete. - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Implementation and recorded runs were AI-assisted with GitHub Copilot. No human-review claim is made. --------- Co-authored-by: kolehma8 <[email protected]> Co-authored-by: Copilot <[email protected]>