mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
f2ac965f61
## Summary Adds a native Blackwell-family (`sm_100f`, including B200/B300) tensor-parallel fused scaled linear cross-entropy implementation. - Fuses the classifier GEMM with NLL/LSE and optional entropy. - Adds a persistent SM100 backward that executes dZ, dX, dW, token-wave scheduling, node-local NVLS, and two-host reduction in one cluster-2 launch. - Supports TP1/2/4/8 node-local execution and two-host hierarchical TP with 1/2/4/8 selected GPUs per host, including TP subgroups of a larger NVSHMEM world. - Uses the shared forward-verified QP put-signal ring for inter-host communication. - Distributes the remote FP32 dX merge over warp 1 from every resident CTA, synchronized with an HBM atomic completion counter. - Adds reusable workspace configuration/query APIs, public API documentation, correctness validators, phase benchmarks, and topology contract tests. ## Code entry points | Layer | Forward | Backward | |---|---|---| | PyTorch autograd | [`LigerFusedLinearScaledCrossEntropyNativeTPFunction.forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L139) | [`LigerFusedLinearScaledCrossEntropyNativeTPFunction.backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L186) | | TVM FFI | [`fused_linear_scaled_cross_entropy_forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L535) | [`fused_linear_scaled_cross_entropy_backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L610) | | SM100 host launcher | [`fused_linear_scaled_cross_entropy_forward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward_sm100.cu#L703) | [`fused_linear_scaled_cross_entropy_backward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward_sm100.cu#L445) | | Device kernel | [`forward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_kernel_sm100.cuh#L26) | [`backward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh#L1109) | | Role implementations | [`forward_gemm_roles_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_roles_sm100.cuh) | [`backward_gemm_mainloop_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh) | ## Multi-host subgroup topology Multi-host setup is relative to the configured TP team rather than the full NVSHMEM world. `nvshmem_team_split_2d` derives: - one local NVLS row containing only the selected GPUs on the current host; - one matching-rank remote column containing the corresponding selected GPU from each host. For a 16-rank world on two eight-GPU hosts, TP group `{0,4,8,12}` becomes local rows `{0,4}` / `{8,12}` and remote pairs `{0,8}` / `{4,12}`. SM100 remote kernels use a parent-team stream barrier followed by a cooperative clustered CUDA launch, so disjoint TP groups do not require world-synchronous kernel entry. Validated concurrently as four TP4 groups `{0,4,8,12}`, `{1,5,9,13}`, `{2,6,10,14}`, and `{3,7,11,15}`: - 28/28 numerical cases pass across the four groups. - Covers entropy, ignored targets, ragged vocabulary, full waves, and two-wave tails. - Cross-rank dX equality remains exact. - Deliberately staggering group entry by 500 ms also passes for forward and backward, confirming subgroup-scoped progress. - Two-host local widths 1, 2, 4, and 8 GPUs per host pass. - `NVSHMEM_IB_ENABLE_IBGDA=1` was forced for the final campaign. Runtime initialization reported `Successfully initialized the transport: IBGDA. It will be used for device-side APIs over IB`, created DCI/RC QPs, and selected each configured `mlx5_*` device. Random BF16, `M=4096`, `H=4096`, global `V=131072`, 10 warmups, and 30 measured iterations: | TP group root | Native forward ms | Native forward TF/GPU | Forward speedup | Native F+B ms | Native F+B TF/GPU | F+B speedup | |---:|---:|---:|---:|---:|---:|---:| | 0 | 1.080 | 1017.9 | 2.82x | 5.193 | 846.9 | 1.69x | | 1 | 1.072 | 1025.8 | 2.85x | 5.166 | 851.3 | 1.69x | | 2 | 1.058 | 1039.5 | 2.86x | 5.128 | 857.6 | 1.69x | | 3 | 1.090 | 1009.2 | 2.78x | 5.175 | 849.9 | 1.68x | Native versus fallback forward maximum absolute difference was `5.45e-4`. ### SM90 regression validation Validated the shared topology, workspace, and binding changes on two eight-GPU H200 hosts with a clean `sm_90a` build: - Complete native core build passed. - 12/12 shared reduction and parent-relative topology tests passed. - 8/8 SM90 FSLCE configuration and online-softmax contract tests passed. - 12/12 standalone SM90 forward GEMM numerical cases passed. - Public FFI world-team and process-group subteam tests passed, including forward, dX, dW, entropy, and CUDA graph replay. - A two-host TP16 world team passed forward, dX, dW, and entropy validation. - Four concurrent strided TP4 groups `{0,4,8,12}`, `{1,5,9,13}`, `{2,6,10,14}`, and `{3,7,11,15}` all passed on H200. - NVSHMEM reported IBGDA as the active device-side InfiniBand transport on all 16 ranks. ## Implementation ### Forward - 2SM UMMA, `M256 x N256 x K64`, cluster `2 x 1`. - Five-stage TMA mainloop. - Four source slots and a default wave width of 64 N256 vocabulary tiles. - Warp 0 performs local reduction, warp 1 runs the matching-rank host ring, warp 2 owns TMA, warp 3 issues UMMA, and warps 4-11 drain TMEM/epilogue work. - The `__global__` entry point is isolated in `forward_gemm_kernel_sm100.cuh` and dispatches those roles explicitly; templated implementations live in `forward_gemm_roles_sm100.cuh`. ### Backward - One persistent 384-thread cluster-2 kernel. - dZ: five-stage `M256 x N256 x K64`. - dX: four-stage 2SM N512 schedule, publishing two M128xN256 communication tiles. - dW: six-stage paired-N schedule with TMA store for wave 0 and TMA reduce-add for later waves. - A single 512-column `Allocator2Sm` TMEM allocation is reused across all phases. - TP16: 1. every CTA's warp 0 performs node-local NVLS reduce-scatter; 2. CTA 0 warp 1 performs the existing inter-host ring transfer; 3. warp 1 from all resident CTAs merges a global-strided portion of the remote FP32 shard; 4. one HBM atomic arrival per CTA releases CTA 0; 5. warp 0 performs node-local all-gather and BF16 scatter. The measured TP16 remote stage is 0.285-0.328 ms: | Component | Latency | |---|---:| | IBRC transfer, 8 MiB/GPU | 0.119-0.156 ms | | All-CTA FP32 merge | 0.160-0.176 ms | | Total remote stage | 0.285-0.328 ms | | In-kernel dW overlap window | ~2.41 ms | | Exposed remote tail | 0 ms | ## Correctness - TP1/2/4/8/16: 35/35 distributed numerical cases pass. - Covers aligned and ragged shapes, entropy, ignored targets, 4K waves, 4097-token boundaries, and 5000-token two-wave tails. - Cross-rank dX equality is exact. - Maximum observed relative error: - dX: 0.262% - dW: 0.498% - Forward native versus Verl maximum absolute difference: `1.09e-3`. - C++ tests: - shared reduction/topology: 11/11 - SM100 backward contracts: 12/12 - `make checkstyle` passes. ## Performance methodology - GPU: NVIDIA B300. - BF16 inputs, FP32 accumulation. - Hidden size: 4096. - Fixed global vocabulary: 131072. - Tokens: 4096, 8192, 16384. - TP: 1, 2, 4, 8, 16. - 10 warmups, 30 measured iterations, median CUDA-event latency. - Verl baseline: repository 512-token chunked TP fallback. - Forward+backward throughput counts four classifier GEMMs: forward logits, backward logits recomputation for dZ, dX, and dW. - Efficiency is per-GPU TFLOP/s relative to TP1 for the same shape. - TP16 uses two eight-GPU hosts and explicit one-HCA-per-local-GPU mapping. ### M=4096 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 2.852 | 1542.0 | 100.0% | 8.489 | 518.1 | 100.0% | 2.98x | | 2 | 1.457 | 1508.9 | 97.9% | 4.653 | 472.6 | 91.2% | 3.19x | | 4 | 0.764 | 1438.3 | 93.3% | 2.856 | 384.9 | 74.3% | 3.74x | | 8 | 0.428 | 1283.6 | 83.2% | 2.134 | 257.6 | 49.7% | 4.98x | | 16 | 0.465 | 591.0 | 38.3% | 2.096 | 131.1 | 25.3% | 4.51x | ### M=4096 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 10.950 | 1606.5 | 100.0% | 29.898 | 588.4 | 100.0% | 2.73x | | 2 | 6.170 | 1425.6 | 88.7% | 15.570 | 564.9 | 96.0% | 2.52x | | 4 | 3.562 | 1234.7 | 76.9% | 8.451 | 520.4 | 88.5% | 2.37x | | 8 | 2.253 | 976.0 | 60.8% | 5.426 | 405.2 | 68.9% | 2.41x | | 16 | 2.171 | 506.5 | 31.5% | 4.472 | 245.9 | 41.8% | 2.06x | ### M=8192 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 5.744 | 1531.4 | 100.0% | 17.107 | 514.2 | 100.0% | 2.98x | | 2 | 2.867 | 1533.8 | 100.2% | 9.376 | 469.1 | 91.2% | 3.27x | | 4 | 1.469 | 1496.9 | 97.7% | 5.638 | 390.1 | 75.9% | 3.84x | | 8 | 0.793 | 1385.8 | 90.5% | 4.185 | 262.7 | 51.1% | 5.28x | | 16 | 0.762 | 721.6 | 47.1% | 4.077 | 134.8 | 26.2% | 5.35x | ### M=8192 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 22.264 | 1580.3 | 100.0% | 58.375 | 602.8 | 100.0% | 2.62x | | 2 | 11.600 | 1516.5 | 96.0% | 31.618 | 556.4 | 92.3% | 2.73x | | 4 | 6.463 | 1360.9 | 86.1% | 17.348 | 507.1 | 84.1% | 2.68x | | 8 | 4.219 | 1042.5 | 66.0% | 10.653 | 412.8 | 68.5% | 2.52x | | 16 | 4.020 | 546.9 | 34.6% | 8.659 | 254.0 | 42.1% | 2.15x | ### M=16384 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 12.645 | 1391.2 | 100.0% | 35.020 | 502.4 | 100.0% | 2.77x | | 2 | 6.399 | 1374.5 | 98.8% | 19.560 | 449.7 | 89.5% | 3.06x | | 4 | 3.206 | 1371.9 | 98.6% | 11.426 | 384.9 | 76.6% | 3.56x | | 8 | 1.674 | 1313.3 | 94.4% | 8.197 | 268.3 | 53.4% | 4.90x | | 16 | 1.449 | 759.0 | 54.6% | 9.809 | 112.1 | 22.3% | 6.77x | ### M=16384 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 50.387 | 1396.5 | 100.0% | 115.214 | 610.8 | 100.0% | 2.29x | | 2 | 23.270 | 1512.0 | 108.3% | 62.505 | 562.9 | 92.2% | 2.69x | | 4 | 12.579 | 1398.5 | 100.1% | 35.498 | 495.6 | 81.1% | 2.82x | | 8 | 8.281 | 1062.1 | 76.1% | 21.188 | 415.2 | 68.0% | 2.56x | | 16 | 7.898 | 556.8 | 39.9% | 16.706 | 263.2 | 43.1% | 2.12x | ## Memory Inputs, weights, and final outputs are excluded. Native reports its exact reusable C++/NVSHMEM forward+backward workspace. Verl reports incremental PyTorch peak memory. ### M=4096 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 1810.40 | 1024.10 | 2720.11 | | 2 | 1202.42 | 512.10 | 1376.11 | | 4 | 898.46 | 256.10 | 704.11 | | 8 | 746.53 | 128.10 | 368.11 | | 16 | 682.67 | 64.10 | 200.11 | ### M=8192 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2004.79 | 1024.19 | 2752.21 | | 2 | 1300.81 | 512.19 | 1408.21 | | 4 | 948.84 | 256.19 | 736.21 | | 8 | 772.91 | 128.19 | 400.21 | | 16 | 709.06 | 64.19 | 232.21 | ### M=16384 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2393.56 | 1024.35 | 2816.40 | | 2 | 1497.58 | 512.35 | 1472.40 | | 4 | 1049.61 | 256.35 | 800.40 | | 8 | 825.68 | 128.35 | 464.40 | | 16 | 761.83 | 64.35 | 296.40 | ## Additional weak-scaling production result With fixed `V_local=65536`, `M=4096`, and `H=4096`: | TP | Native forward TF/GPU | Verl forward TF/GPU | Native F+B TF/GPU | Verl F+B TF/GPU | |---:|---:|---:|---:|---:| | 1 | 1524.5 | 488.6 | 1607.1 | 570.3 | | 2 | 1505.3 | 471.8 | 1424.4 | 560.7 | | 4 | 1513.3 | 459.2 | 1541.9 | 543.2 | | 8 | 1490.8 | 446.8 | 1580.4 | 537.9 | | 16 | 1323.0 | 412.2 | 1434.5 | 533.7 | ## Model-shape campaign The following production model shapes use `M=8192`, random BF16 inputs, 10 warmups, and 30 measured iterations. TP16 spans the same two eight-GPU hosts as the fixed-shape campaign. ### Qwen3-30B-A3B (`H=2048`, `V=151936`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 3.692 | 1380.8 | 100.0% | 15.885 | 320.9 | 100.0% | 4.30x | | 2 | 1.937 | 1315.9 | 95.3% | 9.192 | 277.3 | 86.4% | 4.75x | | 4 | 1.040 | 1225.4 | 88.7% | 5.581 | 228.4 | 71.2% | 5.37x | | 8 | 0.592 | 1076.2 | 77.9% | 4.136 | 154.1 | 48.0% | 6.99x | | 16 | 0.664 | 479.8 | 34.7% | 4.012 | 79.4 | 24.7% | 6.04x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 13.692 | 1489.3 | 100.0% | 47.689 | 427.6 | 100.0% | 3.48x | | 2 | 7.480 | 1363.2 | 91.5% | 26.363 | 386.8 | 90.4% | 3.52x | | 4 | 4.100 | 1243.3 | 83.5% | 14.834 | 343.7 | 80.4% | 3.62x | | 8 | 2.507 | 1016.5 | 68.3% | 9.720 | 262.3 | 61.3% | 3.88x | | 16 | 2.545 | 500.8 | 33.6% | 8.346 | 152.7 | 35.7% | 3.28x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 1975.79 | 1187.19 | 1962.08 | | 2 | 1286.31 | 593.69 | 996.65 | | 4 | 941.84 | 296.94 | 515.33 | | 8 | 769.41 | 152.19 | 276.94 | | 16 | 695.56 | 74.37 | 154.58 | ### Llama 4 (`H=5120`, `V=202048`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 11.727 | 1445.3 | 100.0% | 28.506 | 594.6 | 100.0% | 2.43x | | 2 | 5.795 | 1462.3 | 101.2% | 15.228 | 556.5 | 93.6% | 2.63x | | 4 | 2.880 | 1471.2 | 101.8% | 8.489 | 499.1 | 83.9% | 2.95x | | 8 | 1.479 | 1433.0 | 99.1% | 5.339 | 396.8 | 66.7% | 3.61x | | 16 | 1.103 | 960.3 | 66.4% | 7.578 | 139.8 | 23.5% | 6.87x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 50.180 | 1351.1 | 100.0% | 100.330 | 675.7 | 100.0% | 2.00x | | 2 | 22.677 | 1494.8 | 110.6% | 54.966 | 616.7 | 91.3% | 2.42x | | 4 | 11.785 | 1438.3 | 106.4% | 29.663 | 571.5 | 84.6% | 2.52x | | 8 | 6.939 | 1221.3 | 90.4% | 16.230 | 522.1 | 77.3% | 2.34x | | 16 | 6.026 | 703.2 | 52.0% | 23.633 | 179.3 | 26.5% | 3.92x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2655.29 | 1578.69 | 5015.46 | | 2 | 1626.31 | 792.19 | 2547.99 | | 4 | 1111.84 | 394.81 | 1315.52 | | 8 | 854.41 | 200.19 | 698.15 | | 16 | 756.06 | 98.84 | 389.87 | ### DeepSeek-V3 (`H=7168`, `V=129280`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 10.923 | 1390.0 | 100.0% | 22.352 | 679.2 | 100.0% | 2.05x | | 2 | 5.317 | 1427.6 | 102.7% | 11.923 | 636.7 | 93.7% | 2.24x | | 4 | 2.638 | 1439.0 | 103.5% | 6.943 | 546.7 | 80.5% | 2.63x | | 8 | 1.356 | 1399.8 | 100.7% | 4.556 | 416.5 | 61.3% | 3.36x | | 16 | 1.041 | 911.7 | 65.6% | 4.862 | 195.2 | 28.7% | 4.67x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 43.588 | 1393.3 | 100.0% | 81.751 | 742.9 | 100.0% | 1.88x | | 2 | 22.167 | 1369.9 | 98.3% | 44.494 | 682.4 | 91.9% | 2.01x | | 4 | 12.411 | 1223.3 | 87.8% | 24.582 | 617.6 | 83.1% | 1.98x | | 8 | 7.600 | 998.8 | 71.7% | 13.475 | 563.3 | 75.8% | 1.77x | | 16 | 7.024 | 540.4 | 38.8% | 10.843 | 350.0 | 47.1% | 1.54x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2278.79 | 1010.19 | 4279.46 | | 2 | 1437.81 | 505.19 | 2196.71 | | 4 | 1017.34 | 256.19 | 1156.21 | | 8 | 807.41 | 128.19 | 634.08 | | 16 | 744.56 | 64.19 | 373.04 | ## Validation commands ```bash make checkstyle ``` On B300 (`-DLIGER_CUTE_CUDA_ARCH=100f`): ```bash cmake --build build --target \ libliger_cute_kernels.so \ test_fslce_shared_reduction \ test_fused_scaled_linear_cross_entropy_backward_sm100 ./build/tests/cpp/test_fslce_shared_reduction ./build/tests/cpp/test_fused_scaled_linear_cross_entropy_backward_sm100 ``` Distributed numerical validation: ```bash python -m torch.distributed.run --standalone --nproc_per_node=<1|2|4|8> \ liger_cute_kernels/tests/python/validate_fused_scaled_linear_cross_entropy_backward_sm100.py ``` TP16 uses two eight-GPU hosts with: ```bash export NVSHMEM_REMOTE_TRANSPORT=ibrc export NVSHMEM_IB_ENABLE_IBGDA=1 export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 unset NVSHMEM_HCA_LIST export NVSHMEM_HCA_PE_MAPPING='mlx5_0:1:1,mlx5_2:1:1,mlx5_3:1:1,mlx5_4:1:1,mlx5_5:1:1,mlx5_6:1:1,mlx5_8:1:1,mlx5_9:1:1' ``` --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]>