From 99747c53ba48607395e0160a11a1d11481e3501c Mon Sep 17 00:00:00 2001 From: Jari Kolehmainen Date: Mon, 31 Aug 2026 20:15:17 +0000 Subject: [PATCH] [CUDA] Add Hopper tensor-parallel fused linear cross entropy (#1422) ## Summary Adds Hopper tensor-parallel fused linear scaled cross entropy forward/backward with NVLS, direct-peer, and hierarchical two-host paths. Production choices: - zero-spill cluster-2 M-fast dZ handoff with compact precomputed coefficients; - zero-spill four-stage combined dX+dW with one reused R24 accumulator; - H2048 dX split-K2 when the local vocabulary has an even K64 tile count, otherwise split-K1; H4096 uses split-K1; - N32 FP32 TMA dX publication, N64/S2 dW epilogue, final-only publication fence; - post-compute local/hierarchical FP32-to-BF16 finalization; - device-resident graph-safe launch epochs with one stream-ordered tail barrier; - non-RDC WGMMA translation units with remote NVSHMEM code isolated in RDC sources. ## Tensor-parallel reduction API Reduction transport is now reusable infrastructure under `liger_cute::detail`: - `TpReducePlan` selects the local NVLS or DirectPeer backend on the host. - Kernels receive only the selected compact device view; the full host plan is never passed to device code. - `RequiresRemote` is a launcher template argument, so the local path carries no runtime remote branch or remote fields. - Two-host execution is local NVLS reduce-scatter, explicit `launch_remote_reduce(...)`, then the common local all-gather/scatter. - DirectPeer remains supported behind the same local reduction API. The refactor removes the unused pre-cluster dZ/role framework and old non-cluster hierarchical implementation. Net change: 1,907 fewer lines while preserving the validated kernel schedules. ## Kernel structure ### Warp functions and roles **Forward `forward_gemm_kernel_sm90` (12 warps, cluster-2)** - Warp 0 is the TMA producer for X and weight panels. - Warps 1-3 are reserved/inactive. - Warp groups 1-2 (8 warps) execute the N320 logical GEMM as two N160 accumulator banks, then perform the N80/N40 online-softmax and entropy fold. - Cluster peers share the token-side input while owning adjacent vocabulary work. **Backward dZ handoff (12 warps, cluster-2)** - Warp 3 drives the three-stage TMA pipeline for X/W panels. - Warps 0-2 consume handed-off FP16 logits, apply precomputed row coefficients, target correction, BF16 conversion, and TMA output stores. - Warp groups 1-2 execute M128xN256xK64 WGMMA and hand the rounded FP16 logits through shared memory. - The handoff separates accumulator and epilogue liveness, eliminating the old dZ spill path. **Backward combined dX+dW (12 warps, cluster-2)** - Warp 0 drives the continuous four-stage TMA producer pipeline. - Two communication warps consume the FP32 dX ring, perform local NVLS reduction/publication, and overlap with dW. - Warp groups 1-2 execute dX then dW using one explicitly reused M128xN256 FP32 accumulator bank at R24. - dX uses W multicast and N32 FP32 TMA publication. dW uses dZ multicast, independent X loads, and N64/S2 TMA store/reduce-add. - H2048 uses split-K2 with atomic accumulation into a stream-cleared durable buffer when its K64 tile count is even; odd tile counts and H4096 use split-K1. ### Forward launch sequence 1. `forward_gemm_kernel_sm90`: local GEMM plus online max/sum/entropy/target partials. 2. `forward_split_reduce_kernel_sm90`: reduce local vocabulary splits. 3. Tensor-parallel MAX reduction. 4. `pack_forward_reduction_kernel`: correct/pack sum fields. 5. Tensor-parallel SUM reduction. 6. `finalize_forward_kernel`: produce log-probability, LSE, and optional entropy. ### Backward launch sequence For each token wave: 1. `backward_dz_handoff_wave_kernel_sm90`: logits GEMM, compact dZ math, BF16 dZ workspace. 2. `backward_dx_cluster2_gemm_wave_kernel_sm90`: continuous four-stage dX then dW compute with overlapped local communication. After all waves: 3. Single host: local FP32 reduction result to BF16 scatter. 4. Hierarchical: separate-RDC inter-host reduction, local NVLS broadcast/all-gather, then BF16 scatter. 5. Direct-peer systems use the post-compute fallback finalizer. 6. One stream-ordered tensor-parallel tail barrier prevents staging reuse before peers finish. ### Hoisted versus repeated work **Hoisted** - Backend capability/team selection, workspace reservation, rank mapping, and TMA descriptor construction occur once on the host. - TMA descriptors are prefetched once per CTA role before persistent tile loops. - Shape counts, raster bounds, and split-K selection are computed outside inner K loops. - dZ stores per-row `exp_bias`, `correction_bias`, and `correction_slope`; row coefficients/target are hoisted outside the inner column loop. - dX and dW share one accumulator object and aliased A/B shared-memory arenas. - IB communication and final BF16 materialization remain outside the hot combined WGMMA kernel. **Not hoisted** - Per-wave dZ and combined dX+dW launches. - Per-tile TMA loads, WGMMA issue, accumulator-to-SMEM conversion, and dX publication. - Per-element dZ `exp2`, entropy correction, and target subtraction. - Per-stage NVLS ring synchronization/reduction and split-K2 atomic accumulation. - dW N64 TMA store/reduce-add for each output subtile and wave. ## Validation - `make checkstyle` - clean CMake shared-library build - C++ tests: 8/8 passed - Python WORLD and strided-subteam forward/backward tests: 2/2 passed, including two CUDA Graph replays - backward TP1/2/4/8 selftests: entropy on/off, ragged shapes, multi-wave/full-grid, TPR1/2/4 - poison staging/dZ and delayed communication warp tests - forward TP1/2/4/8 selftests: entropy on/off, ragged, temperature, large split - hierarchical TP16 forward/backward - benchmarked kernel schedule before the transport-only API refactor: 92/92 unique matched points passed in both entropy modes; TP1/2/4/8/16 - post-refactor: checkstyle and focused Python facade regression passed; GPU compilation remains delegated to PR validation ## Benchmark methodology H200 SXM, CUDA 12.9, 8 warmups and 30 measured iterations, p50 shown below. TFLOP/s are per GPU unless labeled aggregate. Liger entropy-off is compared with Verl's no-entropy configuration. Memory factor is `Verl GiB / Liger GiB`; values above 1 mean Liger uses less memory. Allocation-free means no hot-path allocation; Liger's reported footprint includes pre-reserved symmetric communication workspaces. H100 and H200 absolute results are kept separate. ## Overall Liger vs Verl | Pass | Matched points | Geomean speedup | Range | Slower points | Geomean mem V/L | |---|---:|---:|---:|---:|---:| | backward | 46 | 1.54x | 0.95-1.87x | 1 | 0.82x | | forward | 46 | 2.11x | 1.69-3.79x | 0 | 0.47x | Forward is faster at every matched point. Backward is faster at every matched point except M8192/H2048/V131072 at TP16. ## Strong-scaling efficiency Efficiency is per-GPU TFLOP/s at TPn divided by TP1 per-GPU TFLOP/s for the same global workload. | Pass | M | H | V | TP1 TF/GPU | TP2 eff. | TP4 eff. | TP8 eff. | TP16 eff. | |---|---:|---:|---:|---:|---:|---:|---:|---:| | backward | 2048 | 4096 | 131072 | 767.9 | 98.1% | 92.7% | 82.5% | 50.1% | | backward | 8192 | 2048 | 131072 | 694.9 | 91.4% | 82.8% | 68.6% | 37.8% | | backward | 8192 | 4096 | 131072 | 762.8 | 97.8% | 93.7% | 84.0% | 54.5% | | backward | 16384 | 4096 | 131072 | 761.1 | 97.3% | 93.8% | 83.7% | 55.3% | | backward | 2048 | 4096 | 262144 | 783.1 | 98.2% | 97.0% | 91.1% | 66.8% | | backward | 8192 | 2048 | 262144 | 711.3 | 96.7% | 92.3% | 85.0% | 58.2% | | backward | 8192 | 4096 | 262144 | 777.4 | 98.4% | 96.6% | 91.8% | 70.6% | | backward | 16384 | 4096 | 262144 | 775.8 | 98.5% | 95.8% | 91.9% | 71.2% | | forward | 2048 | 4096 | 131072 | 787.7 | 99.9% | 97.7% | 83.5% | 58.4% | | forward | 8192 | 2048 | 131072 | 747.7 | 99.4% | 96.7% | 90.4% | 79.1% | | forward | 8192 | 4096 | 131072 | 789.0 | 99.7% | 98.8% | 95.9% | 88.9% | | forward | 16384 | 4096 | 131072 | 757.4 | 99.7% | 98.3% | 95.0% | 91.4% | | forward | 2048 | 4096 | 262144 | 791.1 | 99.0% | 100.0% | 95.4% | 77.0% | | forward | 8192 | 2048 | 262144 | 747.5 | 99.7% | 99.1% | 95.6% | 86.8% | | forward | 8192 | 4096 | 262144 | 788.3 | 99.9% | 100.1% | 98.3% | 94.1% | | forward | 16384 | 4096 | 262144 | 755.3 | 100.0% | 99.8% | 98.3% | 93.9% | ## Weak-scaling efficiency Global vocabulary grows with TP while local vocabulary remains 65,536. Efficiency is per-GPU TFLOP/s relative to TP1. | Pass | H | TP1 TF/GPU | TP2 eff. | TP4 eff. | TP8 eff. | TP16 eff. | |---|---:|---:|---:|---:|---:|---:| | backward | 2048 | 661.4 | 96.1% | 99.2% | 99.2% | 94.7% | | backward | 4096 | 757.6 | 98.4% | 99.1% | 98.8% | 95.8% | | forward | 2048 | 748.2 | 99.4% | 99.0% | 98.5% | 97.4% | | forward | 4096 | 788.5 | 99.7% | 100.1% | 99.3% | 98.4% | ## Full Liger vs Verl results | Pass | M | H | V | TP | Liger TF/GPU | Verl TF/GPU | Speedup | Liger GiB | Verl GiB | Mem V/L | |---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| | backward | 8192 | 2048 | 65536 | 1 | 661.4 | 375.2 | 1.76x | 2.26 | 1.15 | 0.51x | | backward | 8192 | 4096 | 65536 | 1 | 757.6 | 507.3 | 1.49x | 2.82 | 1.99 | 0.71x | | backward | 2048 | 4096 | 131072 | 1 | 767.9 | 445.2 | 1.72x | 3.85 | 3.14 | 0.82x | | backward | 2048 | 4096 | 131072 | 2 | 753.3 | 436.4 | 1.73x | 2.72 | 1.65 | 0.61x | | backward | 2048 | 4096 | 131072 | 4 | 711.6 | 385.0 | 1.85x | 2.16 | 0.88 | 0.41x | | backward | 2048 | 4096 | 131072 | 8 | 633.3 | 362.6 | 1.75x | 1.88 | 0.52 | 0.28x | | backward | 2048 | 4096 | 131072 | 16 | 385.0 | 311.3 | 1.24x | 0.17 | 0.33 | 1.94x | | backward | 8192 | 2048 | 131072 | 1 | 694.9 | 378.9 | 1.83x | 2.88 | 1.89 | 0.66x | | backward | 8192 | 2048 | 131072 | 2 | 635.4 | 372.0 | 1.71x | 2.26 | 1.15 | 0.51x | | backward | 8192 | 2048 | 131072 | 4 | 575.2 | 329.3 | 1.75x | 1.94 | 0.71 | 0.37x | | backward | 8192 | 2048 | 131072 | 8 | 476.9 | 313.8 | 1.52x | 1.79 | 0.56 | 0.31x | | backward | 8192 | 2048 | 131072 | 16 | 262.4 | 277.1 | 0.95x | 0.14 | 0.49 | 3.50x | | backward | 8192 | 4096 | 131072 | 1 | 762.8 | 510.6 | 1.49x | 3.94 | 3.48 | 0.88x | | backward | 8192 | 4096 | 131072 | 2 | 745.7 | 501.2 | 1.49x | 2.82 | 1.99 | 0.71x | | backward | 8192 | 4096 | 131072 | 4 | 714.7 | 446.2 | 1.60x | 2.26 | 1.18 | 0.52x | | backward | 8192 | 4096 | 131072 | 8 | 640.6 | 431.4 | 1.48x | 1.97 | 0.84 | 0.43x | | backward | 8192 | 4096 | 131072 | 16 | 415.7 | 387.3 | 1.07x | 0.27 | 0.68 | 2.52x | | backward | 16384 | 4096 | 131072 | 1 | 761.1 | 519.8 | 1.46x | 4.07 | 3.93 | 0.97x | | backward | 16384 | 4096 | 131072 | 2 | 740.2 | 510.8 | 1.45x | 2.94 | 2.46 | 0.84x | | backward | 16384 | 4096 | 131072 | 4 | 714.1 | 458.3 | 1.56x | 2.38 | 1.58 | 0.66x | | backward | 16384 | 4096 | 131072 | 8 | 637.1 | 444.4 | 1.43x | 2.10 | 1.28 | 0.61x | | backward | 16384 | 4096 | 131072 | 16 | 421.0 | 404.6 | 1.04x | 0.39 | 1.13 | 2.90x | | backward | 2048 | 4096 | 262144 | 1 | 783.1 | 449.4 | 1.74x | 6.10 | 6.14 | 1.01x | | backward | 2048 | 4096 | 262144 | 2 | 768.8 | 444.1 | 1.73x | 3.85 | 3.14 | 0.82x | | backward | 2048 | 4096 | 262144 | 4 | 759.7 | 435.4 | 1.74x | 2.72 | 1.65 | 0.61x | | backward | 2048 | 4096 | 262144 | 8 | 713.6 | 383.2 | 1.86x | 2.16 | 0.88 | 0.41x | | backward | 2048 | 4096 | 262144 | 16 | 523.4 | 351.2 | 1.49x | 0.31 | 0.52 | 1.68x | | backward | 8192 | 2048 | 262144 | 1 | 711.3 | 381.3 | 1.87x | 4.13 | 3.35 | 0.81x | | backward | 8192 | 2048 | 262144 | 2 | 687.8 | 376.7 | 1.83x | 2.88 | 1.89 | 0.66x | | backward | 8192 | 2048 | 262144 | 4 | 656.4 | 371.3 | 1.77x | 2.26 | 1.15 | 0.51x | | backward | 8192 | 2048 | 262144 | 8 | 604.6 | 329.0 | 1.84x | 1.94 | 0.71 | 0.37x | | backward | 8192 | 2048 | 262144 | 16 | 414.0 | 306.1 | 1.35x | 0.22 | 0.56 | 2.55x | | backward | 8192 | 4096 | 262144 | 1 | 777.4 | 512.1 | 1.52x | 6.19 | 6.45 | 1.04x | | backward | 8192 | 4096 | 262144 | 2 | 765.1 | 505.7 | 1.51x | 3.94 | 3.48 | 0.88x | | backward | 8192 | 4096 | 262144 | 4 | 751.1 | 500.6 | 1.50x | 2.82 | 1.99 | 0.71x | | backward | 8192 | 4096 | 262144 | 8 | 714.0 | 447.2 | 1.60x | 2.26 | 1.18 | 0.52x | | backward | 8192 | 4096 | 262144 | 16 | 548.6 | 421.2 | 1.30x | 0.41 | 0.84 | 2.05x | | backward | 16384 | 4096 | 262144 | 1 | 775.8 | 520.5 | 1.49x | 6.32 | 6.87 | 1.09x | | backward | 16384 | 4096 | 262144 | 2 | 763.9 | 516.4 | 1.48x | 4.07 | 3.93 | 0.97x | | backward | 16384 | 4096 | 262144 | 4 | 743.2 | 509.9 | 1.46x | 2.94 | 2.46 | 0.84x | | backward | 16384 | 4096 | 262144 | 8 | 713.2 | 458.0 | 1.56x | 2.38 | 1.58 | 0.66x | | backward | 16384 | 4096 | 262144 | 16 | 552.6 | 429.2 | 1.29x | 0.53 | 1.28 | 2.42x | | backward | 8192 | 2048 | 524288 | 8 | 655.9 | 369.7 | 1.77x | 2.26 | 1.15 | 0.51x | | backward | 8192 | 4096 | 524288 | 8 | 748.4 | 499.0 | 1.50x | 2.82 | 1.99 | 0.71x | | backward | 8192 | 2048 | 1048576 | 16 | 626.5 | 366.9 | 1.71x | 0.69 | 1.15 | 1.67x | | backward | 8192 | 4096 | 1048576 | 16 | 726.1 | 494.5 | 1.47x | 1.25 | 1.99 | 1.59x | | forward | 8192 | 2048 | 65536 | 1 | 748.2 | 361.7 | 2.07x | 0.28 | 0.29 | 1.04x | | forward | 8192 | 4096 | 65536 | 1 | 788.5 | 417.4 | 1.89x | 0.56 | 0.57 | 1.02x | | forward | 2048 | 4096 | 131072 | 1 | 787.7 | 460.7 | 1.71x | 1.02 | 1.02 | 1.00x | | forward | 2048 | 4096 | 131072 | 2 | 786.9 | 439.2 | 1.79x | 0.98 | 0.52 | 0.53x | | forward | 2048 | 4096 | 131072 | 4 | 769.3 | 408.7 | 1.88x | 0.73 | 0.27 | 0.37x | | forward | 2048 | 4096 | 131072 | 8 | 657.7 | 351.8 | 1.87x | 0.60 | 0.14 | 0.23x | | forward | 2048 | 4096 | 131072 | 16 | 460.3 | 253.4 | 1.82x | 0.46 | 0.08 | 0.17x | | forward | 8192 | 2048 | 131072 | 1 | 747.7 | 226.1 | 3.31x | 0.53 | 0.55 | 1.04x | | forward | 8192 | 2048 | 131072 | 2 | 743.4 | 199.3 | 3.73x | 0.74 | 0.29 | 0.39x | | forward | 8192 | 2048 | 131072 | 4 | 722.7 | 343.3 | 2.11x | 0.62 | 0.16 | 0.26x | | forward | 8192 | 2048 | 131072 | 8 | 676.0 | 321.2 | 2.10x | 0.56 | 0.10 | 0.18x | | forward | 8192 | 2048 | 131072 | 16 | 591.1 | 251.8 | 2.35x | 0.45 | 0.06 | 0.13x | | forward | 8192 | 4096 | 131072 | 1 | 789.0 | 224.1 | 3.52x | 1.06 | 1.08 | 1.02x | | forward | 8192 | 4096 | 131072 | 2 | 786.3 | 411.8 | 1.91x | 1.03 | 0.57 | 0.55x | | forward | 8192 | 4096 | 131072 | 4 | 779.4 | 401.0 | 1.94x | 0.78 | 0.32 | 0.41x | | forward | 8192 | 4096 | 131072 | 8 | 756.8 | 388.6 | 1.95x | 0.65 | 0.19 | 0.29x | | forward | 8192 | 4096 | 131072 | 16 | 701.1 | 333.9 | 2.10x | 0.51 | 0.13 | 0.25x | | forward | 16384 | 4096 | 131072 | 1 | 757.4 | 372.8 | 2.03x | 1.13 | 1.16 | 1.03x | | forward | 16384 | 4096 | 131072 | 2 | 754.9 | 365.0 | 2.07x | 1.09 | 0.64 | 0.59x | | forward | 16384 | 4096 | 131072 | 4 | 744.7 | 358.3 | 2.08x | 0.84 | 0.38 | 0.45x | | forward | 16384 | 4096 | 131072 | 8 | 719.9 | 331.1 | 2.17x | 0.72 | 0.25 | 0.35x | | forward | 16384 | 4096 | 131072 | 16 | 692.2 | 327.0 | 2.12x | 0.57 | 0.19 | 0.33x | | forward | 2048 | 4096 | 262144 | 1 | 791.1 | 468.7 | 1.69x | 2.02 | 2.02 | 1.00x | | forward | 2048 | 4096 | 262144 | 2 | 783.1 | 422.3 | 1.85x | 1.48 | 1.02 | 0.69x | | forward | 2048 | 4096 | 262144 | 4 | 790.9 | 437.6 | 1.81x | 0.98 | 0.52 | 0.53x | | forward | 2048 | 4096 | 262144 | 8 | 754.4 | 404.2 | 1.87x | 0.73 | 0.27 | 0.37x | | forward | 2048 | 4096 | 262144 | 16 | 609.5 | 327.5 | 1.86x | 0.52 | 0.14 | 0.27x | | forward | 8192 | 2048 | 262144 | 1 | 747.5 | 372.3 | 2.01x | 1.03 | 1.06 | 1.03x | | forward | 8192 | 2048 | 262144 | 2 | 745.5 | 360.5 | 2.07x | 0.99 | 0.55 | 0.56x | | forward | 8192 | 2048 | 262144 | 4 | 740.7 | 353.1 | 2.10x | 0.74 | 0.29 | 0.39x | | forward | 8192 | 2048 | 262144 | 8 | 714.4 | 341.1 | 2.09x | 0.62 | 0.16 | 0.26x | | forward | 8192 | 2048 | 262144 | 16 | 648.8 | 309.6 | 2.10x | 0.48 | 0.10 | 0.21x | | forward | 8192 | 4096 | 262144 | 1 | 788.3 | 208.0 | 3.79x | 2.06 | 2.09 | 1.01x | | forward | 8192 | 4096 | 262144 | 2 | 787.6 | 233.2 | 3.38x | 1.53 | 1.08 | 0.71x | | forward | 8192 | 4096 | 262144 | 4 | 789.0 | 411.4 | 1.92x | 1.03 | 0.57 | 0.55x | | forward | 8192 | 4096 | 262144 | 8 | 775.0 | 401.4 | 1.93x | 0.78 | 0.32 | 0.41x | | forward | 8192 | 4096 | 262144 | 16 | 741.9 | 379.7 | 1.95x | 0.57 | 0.19 | 0.33x | | forward | 16384 | 4096 | 262144 | 1 | 755.3 | 374.1 | 2.02x | 2.13 | 2.19 | 1.03x | | forward | 16384 | 4096 | 262144 | 2 | 755.0 | 367.4 | 2.05x | 1.59 | 1.16 | 0.73x | | forward | 16384 | 4096 | 262144 | 4 | 753.8 | 364.0 | 2.07x | 1.09 | 0.64 | 0.59x | | forward | 16384 | 4096 | 262144 | 8 | 742.1 | 357.1 | 2.08x | 0.84 | 0.38 | 0.45x | | forward | 16384 | 4096 | 262144 | 16 | 709.3 | 344.8 | 2.06x | 0.64 | 0.25 | 0.39x | | forward | 8192 | 2048 | 524288 | 8 | 736.9 | 351.5 | 2.10x | 0.74 | 0.29 | 0.39x | | forward | 8192 | 4096 | 524288 | 8 | 782.6 | 410.2 | 1.91x | 1.03 | 0.57 | 0.55x | | forward | 8192 | 2048 | 1048576 | 16 | 729.1 | 347.2 | 2.10x | 0.66 | 0.29 | 0.44x | | forward | 8192 | 4096 | 1048576 | 16 | 776.0 | 407.4 | 1.90x | 0.95 | 0.57 | 0.60x | ## Liger entropy-on results | Pass | M | H | V | TP | p50 ms | TF/GPU | Aggregate TF | GiB | Backend | split-K | |---|---:|---:|---:|---:|---:|---:|---:|---:|---|---:| | backward | 2048 | 4096 | 131072 | 1 | 8.589 | 768.1 | 768.1 | 3.85 | nvls | 1 | | backward | 2048 | 4096 | 131072 | 2 | 4.395 | 750.4 | 1500.9 | 2.72 | nvls | 1 | | backward | 2048 | 4096 | 131072 | 4 | 2.311 | 713.7 | 2854.6 | 2.16 | nvls | 1 | | backward | 2048 | 4096 | 131072 | 8 | 1.291 | 638.5 | 5108.2 | 1.88 | nvls | 1 | | backward | 2048 | 4096 | 131072 | 16 | 1.066 | 386.9 | 6190.2 | 0.17 | hierarchical | 1 | | backward | 2048 | 4096 | 262144 | 1 | 16.842 | 783.4 | 783.4 | 6.10 | nvls | 1 | | backward | 2048 | 4096 | 262144 | 2 | 8.570 | 769.8 | 1539.6 | 3.85 | nvls | 1 | | backward | 2048 | 4096 | 262144 | 4 | 4.366 | 755.6 | 3022.2 | 2.72 | nvls | 1 | | backward | 2048 | 4096 | 262144 | 8 | 2.306 | 715.1 | 5720.5 | 2.16 | nvls | 1 | | backward | 2048 | 4096 | 262144 | 16 | 1.574 | 523.8 | 8380.4 | 0.31 | hierarchical | 1 | | backward | 8192 | 2048 | 65536 | 1 | 10.008 | 659.2 | 659.2 | 2.26 | nvls | 2 | | backward | 8192 | 2048 | 131072 | 1 | 19.086 | 691.3 | 691.3 | 2.88 | nvls | 2 | | backward | 8192 | 2048 | 131072 | 2 | 10.421 | 633.0 | 1266.1 | 2.26 | nvls | 2 | | backward | 8192 | 2048 | 131072 | 4 | 5.759 | 572.8 | 2291.2 | 1.94 | nvls | 2 | | backward | 8192 | 2048 | 131072 | 8 | 3.467 | 475.7 | 3805.4 | 1.79 | nvls | 2 | | backward | 8192 | 2048 | 131072 | 16 | 3.141 | 262.6 | 4201.1 | 0.14 | hierarchical | 2 | | backward | 8192 | 2048 | 262144 | 1 | 37.342 | 706.7 | 706.7 | 4.13 | nvls | 2 | | backward | 8192 | 2048 | 262144 | 2 | 19.290 | 684.0 | 1368.0 | 2.88 | nvls | 2 | | backward | 8192 | 2048 | 262144 | 4 | 10.105 | 652.8 | 2611.3 | 2.26 | nvls | 2 | | backward | 8192 | 2048 | 262144 | 8 | 5.454 | 604.8 | 4838.6 | 1.94 | nvls | 2 | | backward | 8192 | 2048 | 262144 | 16 | 3.999 | 412.4 | 6598.0 | 0.22 | hierarchical | 2 | | backward | 8192 | 2048 | 524288 | 8 | 10.070 | 655.1 | 5241.0 | 2.26 | nvls | 2 | | backward | 8192 | 2048 | 1048576 | 16 | 10.573 | 624.0 | 9983.4 | 0.69 | hierarchical | 2 | | backward | 8192 | 4096 | 65536 | 1 | 17.420 | 757.4 | 757.4 | 2.82 | nvls | 1 | | backward | 8192 | 4096 | 131072 | 1 | 34.601 | 762.6 | 762.6 | 3.94 | nvls | 1 | | backward | 8192 | 4096 | 131072 | 2 | 17.774 | 742.3 | 1484.7 | 2.82 | nvls | 1 | | backward | 8192 | 4096 | 131072 | 4 | 9.214 | 716.0 | 2863.9 | 2.26 | nvls | 1 | | backward | 8192 | 4096 | 131072 | 8 | 5.153 | 640.2 | 5121.4 | 1.97 | nvls | 1 | | backward | 8192 | 4096 | 131072 | 16 | 3.963 | 416.2 | 6658.9 | 0.27 | hierarchical | 1 | | backward | 8192 | 4096 | 262144 | 1 | 67.858 | 777.8 | 777.8 | 6.19 | nvls | 1 | | backward | 8192 | 4096 | 262144 | 2 | 34.470 | 765.5 | 1531.1 | 3.94 | nvls | 1 | | backward | 8192 | 4096 | 262144 | 4 | 17.560 | 751.4 | 3005.5 | 2.82 | nvls | 1 | | backward | 8192 | 4096 | 262144 | 8 | 9.209 | 716.4 | 5731.1 | 2.26 | nvls | 1 | | backward | 8192 | 4096 | 262144 | 16 | 6.014 | 548.5 | 8775.6 | 0.41 | hierarchical | 1 | | backward | 8192 | 4096 | 524288 | 8 | 17.647 | 747.7 | 5981.2 | 2.82 | nvls | 1 | | backward | 8192 | 4096 | 1048576 | 16 | 18.181 | 725.7 | 11611.2 | 1.25 | hierarchical | 1 | | backward | 16384 | 4096 | 131072 | 1 | 69.235 | 762.3 | 762.3 | 4.07 | nvls | 1 | | backward | 16384 | 4096 | 131072 | 2 | 35.679 | 739.6 | 1479.2 | 2.94 | nvls | 1 | | backward | 16384 | 4096 | 131072 | 4 | 18.499 | 713.2 | 2853.0 | 2.38 | nvls | 1 | | backward | 16384 | 4096 | 131072 | 8 | 10.382 | 635.4 | 5083.3 | 2.10 | nvls | 1 | | backward | 16384 | 4096 | 131072 | 16 | 7.832 | 421.1 | 6738.3 | 0.39 | hierarchical | 1 | | backward | 16384 | 4096 | 262144 | 1 | 136.168 | 775.2 | 775.2 | 6.32 | nvls | 1 | | backward | 16384 | 4096 | 262144 | 2 | 69.172 | 763.0 | 1526.0 | 4.07 | nvls | 1 | | backward | 16384 | 4096 | 262144 | 4 | 35.475 | 743.9 | 2975.4 | 2.94 | nvls | 1 | | backward | 16384 | 4096 | 262144 | 8 | 18.444 | 715.4 | 5722.9 | 2.38 | nvls | 1 | | backward | 16384 | 4096 | 262144 | 16 | 11.933 | 552.8 | 8845.5 | 0.53 | hierarchical | 1 | | forward | 2048 | 4096 | 131072 | 1 | 2.842 | 773.9 | 773.9 | 1.02 | nvls | | | forward | 2048 | 4096 | 131072 | 2 | 1.392 | 789.8 | 1579.6 | 0.98 | nvls | | | forward | 2048 | 4096 | 131072 | 4 | 0.720 | 763.8 | 3055.3 | 0.73 | nvls | | | forward | 2048 | 4096 | 131072 | 8 | 0.424 | 648.0 | 5183.8 | 0.60 | nvls | | | forward | 2048 | 4096 | 131072 | 16 | 0.303 | 453.9 | 7262.3 | 0.46 | hierarchical | | | forward | 2048 | 4096 | 262144 | 1 | 5.612 | 783.6 | 783.6 | 2.02 | nvls | | | forward | 2048 | 4096 | 262144 | 2 | 2.820 | 779.7 | 1559.3 | 1.48 | nvls | | | forward | 2048 | 4096 | 262144 | 4 | 1.395 | 788.4 | 3153.5 | 0.98 | nvls | | | forward | 2048 | 4096 | 262144 | 8 | 0.736 | 746.5 | 5972.1 | 0.73 | nvls | | | forward | 2048 | 4096 | 262144 | 16 | 0.456 | 603.4 | 9653.6 | 0.52 | hierarchical | | | forward | 8192 | 2048 | 65536 | 1 | 2.957 | 743.6 | 743.6 | 0.28 | nvls | | | forward | 8192 | 2048 | 131072 | 1 | 5.899 | 745.5 | 745.5 | 0.53 | nvls | | | forward | 8192 | 2048 | 131072 | 2 | 2.973 | 739.7 | 1479.4 | 0.74 | nvls | | | forward | 8192 | 2048 | 131072 | 4 | 1.533 | 717.4 | 2869.6 | 0.62 | nvls | | | forward | 8192 | 2048 | 131072 | 8 | 0.815 | 674.2 | 5393.4 | 0.56 | nvls | | | forward | 8192 | 2048 | 131072 | 16 | 0.473 | 580.9 | 9294.3 | 0.45 | hierarchical | | | forward | 8192 | 2048 | 262144 | 1 | 11.777 | 746.9 | 746.9 | 1.03 | nvls | | | forward | 8192 | 2048 | 262144 | 2 | 5.918 | 743.1 | 1486.2 | 0.99 | nvls | | | forward | 8192 | 2048 | 262144 | 4 | 2.988 | 736.0 | 2944.0 | 0.74 | nvls | | | forward | 8192 | 2048 | 262144 | 8 | 1.548 | 710.5 | 5683.6 | 0.62 | nvls | | | forward | 8192 | 2048 | 262144 | 16 | 0.854 | 643.8 | 10300.6 | 0.48 | hierarchical | | | forward | 8192 | 2048 | 524288 | 8 | 3.002 | 732.6 | 5860.9 | 0.74 | nvls | | | forward | 8192 | 2048 | 1048576 | 16 | 3.045 | 722.1 | 11553.5 | 0.66 | hierarchical | | | forward | 8192 | 4096 | 65536 | 1 | 5.596 | 785.9 | 785.9 | 0.56 | nvls | | | forward | 8192 | 4096 | 131072 | 1 | 11.182 | 786.7 | 786.7 | 1.06 | nvls | | | forward | 8192 | 4096 | 131072 | 2 | 5.575 | 788.9 | 1577.7 | 1.03 | nvls | | | forward | 8192 | 4096 | 131072 | 4 | 2.828 | 777.7 | 3110.7 | 0.78 | nvls | | | forward | 8192 | 4096 | 131072 | 8 | 1.458 | 754.1 | 6033.1 | 0.65 | nvls | | | forward | 8192 | 4096 | 131072 | 16 | 0.789 | 697.0 | 11151.2 | 0.51 | hierarchical | | | forward | 8192 | 4096 | 262144 | 1 | 22.375 | 786.2 | 786.2 | 2.06 | nvls | | | forward | 8192 | 4096 | 262144 | 2 | 11.141 | 789.5 | 1579.1 | 1.53 | nvls | | | forward | 8192 | 4096 | 262144 | 4 | 5.594 | 786.2 | 3144.9 | 1.03 | nvls | | | forward | 8192 | 4096 | 262144 | 8 | 2.833 | 776.2 | 6209.9 | 0.78 | nvls | | | forward | 8192 | 4096 | 262144 | 16 | 1.492 | 737.0 | 11791.4 | 0.57 | hierarchical | | | forward | 8192 | 4096 | 524288 | 8 | 5.645 | 779.2 | 6233.2 | 1.03 | nvls | | | forward | 8192 | 4096 | 1048576 | 16 | 5.677 | 774.7 | 12395.4 | 0.95 | hierarchical | | | forward | 16384 | 4096 | 131072 | 1 | 23.312 | 754.6 | 754.6 | 1.13 | nvls | | | forward | 16384 | 4096 | 131072 | 2 | 11.685 | 752.8 | 1505.6 | 1.09 | nvls | | | forward | 16384 | 4096 | 131072 | 4 | 5.932 | 741.4 | 2965.4 | 0.84 | nvls | | | forward | 16384 | 4096 | 131072 | 8 | 3.067 | 716.9 | 5735.6 | 0.72 | nvls | | | forward | 16384 | 4096 | 131072 | 16 | 1.600 | 687.0 | 10992.6 | 0.57 | hierarchical | | | forward | 16384 | 4096 | 262144 | 1 | 46.752 | 752.6 | 752.6 | 2.13 | nvls | | | forward | 16384 | 4096 | 262144 | 2 | 23.374 | 752.6 | 1505.3 | 1.59 | nvls | | | forward | 16384 | 4096 | 262144 | 4 | 11.702 | 751.7 | 3006.6 | 1.09 | nvls | | | forward | 16384 | 4096 | 262144 | 8 | 5.946 | 739.6 | 5917.2 | 0.84 | nvls | | | forward | 16384 | 4096 | 262144 | 16 | 3.116 | 705.7 | 11291.1 | 0.64 | hierarchical | | ## Resource results - compact dZ handoff: 168 registers, zero spills, 11 barriers; dZ-only p50 11.839 ms entropy-off and 11.898 ms entropy-on in the isolated validation build. - combined split-K1 and split-K2: 168 registers, zero spills, 7 barriers, 230,400 B dynamic shared memory; TPR1/2/4 validated. - dX/dW mainloops are within 10% of tuned cuBLASLt; dX exceeds cuBLAS on multicast-friendly shapes and dW is approximately 4% behind. --------- Co-authored-by: jkolehma Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- liger_cute_kernels/CMakeLists.txt | 22 +- liger_cute_kernels/README.md | 35 + liger_cute_kernels/csrc/core/CMakeLists.txt | 107 +- .../liger_cute/detail/local_reduce.cuh | 154 ++ .../core/include/liger_cute/detail/nvls.cuh | 304 +++ .../include/liger_cute/detail/tp_reduce.cuh | 111 + .../csrc/core/src/detail/remote_all_reduce.cu | 84 + .../core/src/detail/remote_all_reduce.cuh | 26 + .../csrc/core/src/detail/tp_reduce.cu | 357 +++ .../backward.cuh | 47 + .../backward_gemm_mainloop_sm90.cuh | 1907 +++++++++++++++++ .../backward_gemm_sm90.cuh | 381 ++++ .../backward_reduce_sm90.cu | 753 +++++++ .../backward_reduce_sm90.cuh | 51 + .../backward_sm90_selftest.cu | 701 ++++++ .../config.cuh | 62 + .../dx_reduce.cuh | 240 +++ .../forward_gemm_mainloop_sm90.cuh | 794 +++++++ .../forward_gemm_sm90.cuh | 523 +++++ .../forward_gemm_sm90_selftest.cu | 414 ++++ .../forward_reduce.cu | 132 ++ .../forward_reduce.cuh | 80 + .../forward_tp_sm90_selftest.cu | 427 ++++ ...ed_linear_scaled_cross_entropy_backward.cu | 542 +++++ ...sed_linear_scaled_cross_entropy_forward.cu | 348 +++ .../fused_scaled_linear_cross_entropy.cuh | 10 + .../gemm_sm90.cuh | 506 +++++ .../online_softmax.cuh | 118 + .../fused_scaled_linear_cross_entropy/state.h | 11 + .../tma_store_sm90.cuh | 36 + .../workspace.cu | 429 ++++ .../workspace.cuh | 147 ++ liger_cute_kernels/csrc/core/src/nvshmem.cu | 3 + .../liger_cute_kernels/tvm_ffi.py | 120 ++ .../liger_cute_kernels/tvm_ffi_bindings.cpp | 263 +++ liger_cute_kernels/test/test_fslce_teams.py | 282 +++ liger_cute_kernels/tests/cpp/CMakeLists.txt | 14 + .../test_fused_scaled_linear_cross_entropy.cu | 195 ++ 38 files changed, 10731 insertions(+), 5 deletions(-) create mode 100644 liger_cute_kernels/csrc/core/include/liger_cute/detail/local_reduce.cuh create mode 100644 liger_cute_kernels/csrc/core/include/liger_cute/detail/nvls.cuh create mode 100644 liger_cute_kernels/csrc/core/include/liger_cute/detail/tp_reduce.cuh create mode 100644 liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cu create mode 100644 liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cuh create mode 100644 liger_cute_kernels/csrc/core/src/detail/tp_reduce.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_sm90_selftest.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/config.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/dx_reduce.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_mainloop_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90_selftest.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_tp_sm90_selftest.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_scaled_linear_cross_entropy.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/gemm_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/online_softmax.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/state.h create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/tma_store_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cu create mode 100644 liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cuh create mode 100644 liger_cute_kernels/test/test_fslce_teams.py create mode 100644 liger_cute_kernels/tests/cpp/test_fused_scaled_linear_cross_entropy.cu diff --git a/liger_cute_kernels/CMakeLists.txt b/liger_cute_kernels/CMakeLists.txt index c7887ff3..be6ca1f9 100644 --- a/liger_cute_kernels/CMakeLists.txt +++ b/liger_cute_kernels/CMakeLists.txt @@ -50,6 +50,21 @@ option(LIGER_CUTE_TESTS_ONLY # ---------- Common CUDA / device dependencies (needed by core AND bindings) -- find_package(CUDAToolkit REQUIRED) +if(NOT LIGER_CUTE_TESTS_ONLY) + if(NOT DEFINED NCCL_HOME AND DEFINED ENV{NCCL_HOME}) + set(NCCL_HOME "$ENV{NCCL_HOME}") + endif() + find_path(NCCL_INCLUDE_DIR NAMES nccl.h + HINTS "${NCCL_HOME}/include" + PATHS /usr/include /usr/local/cuda/include) + find_library(NCCL_LIBRARY NAMES nccl + HINTS "${NCCL_HOME}/lib" "${NCCL_HOME}/lib64" + PATHS /usr/lib /usr/lib64 /usr/lib/x86_64-linux-gnu) + if(NOT NCCL_INCLUDE_DIR OR NOT NCCL_LIBRARY) + message(FATAL_ERROR "NCCL not found (set NCCL_HOME=)") + endif() +endif() + # The 'a' suffix enables architecture-accelerated features (WGMMA/TMA/multicast # on sm_90a; tcgen05/UMMA on sm_100a). Override with -DLIGER_CUTE_CUDA_ARCH=100a # for a Blackwell build. Strip any gencode flags a parent/toolchain may have @@ -71,10 +86,9 @@ endif() # Disable device-side assertions to avoid extern __assert_fail calls that cause # ptxas to serialize WGMMA instructions. add_compile_definitions(NDEBUG) -# NVSHMEM_ENABLE_ALL_DEVICE_INLINING is REQUIRED globally — see LigerCommKernels -# root CMakeLists for the full rationale (nvshmem.h device helpers reference -# device-state symbols gated on __CUDACC_RDC__; inlining makes non-RDC host-only -# TUs build cleanly). +# NVSHMEM_ENABLE_ALL_DEVICE_INLINING is the default for non-RDC host-only TUs: +# nvshmem.h device helpers otherwise reference state gated on __CUDACC_RDC__. +# RDC collective TUs explicitly undefine it when they need library device code. add_compile_definitions(NVSHMEM_ENABLE_ALL_DEVICE_INLINING=ON) # ── NVSHMEM ────────────────────────────────────────────────────────────────── diff --git a/liger_cute_kernels/README.md b/liger_cute_kernels/README.md index 79d986d6..a132d078 100644 --- a/liger_cute_kernels/README.md +++ b/liger_cute_kernels/README.md @@ -48,6 +48,41 @@ global batch size 512, and sequence length 4,096. Left: cross-entropy loss. Right: per-step wall time. The paper excludes the initial CUDA Graph build from its aggregate step-time statistics. +## Fused linear scaled cross entropy + +The Hopper native core also provides +`fused_linear_scaled_cross_entropy_forward` and +`fused_linear_scaled_cross_entropy_backward`. The backward path uses a +three-stage cluster-2 dZ handoff followed by a four-stage combined dX+dW +cluster kernel. dX uses split-K=2 only for hidden size 2,048; hidden size 4,096 +uses split-K=1. Direct-peer and hierarchical two-host transports remain +available when the single-host NVLS path is not selected. + +The `liger_cute_kernels.tvm_ffi` facade exposes configuration plus forward and +backward entry points. Forward accepts the caller's tensor-parallel +`ncclComm_t` as an integer handle; backward accepts the NVSHMEM team configured +for the same tensor-parallel ranks. Both calls take inverse temperature. +Capacities for tokens, hidden size, local vocabulary, and reduction grouping +must be configured before capture or execution. + +The facade owns NCCL communicator bootstrap helpers: +`nccl_get_unique_id`, `nccl_comm_init_rank`, and `nccl_comm_destroy`. Broadcast +the returned CPU `uint8` unique-ID tensor to the tensor-parallel ranks, create +one communicator per rank, and pass its integer handle to forward. + +Hierarchical NVLS+remote backward currently requires the configured +tensor-parallel team to cover the full NVSHMEM world. Single-host NVLS +subgroups remain supported; a multi-host subgroup that is only part of a +larger world fails explicitly rather than communicating with nonmembers. + +Tensor-parallel reduction transport is implemented in +`liger_cute::detail`. Host setup selects either the NVLS or DirectPeer local +backend and passes only that backend's compact device view to the kernel. +Two-host execution is composed as local NVLS reduce-scatter, an explicit +RDC-compiled remote reduction follow-up, and the same local all-gather/scatter +used by single-host execution. Whether the remote follow-up is required is a +launcher template parameter, not a runtime kernel argument. + ## Package architecture Native CUDA build for the MoE port (from `LigerCommKernels`). The lck wheel diff --git a/liger_cute_kernels/csrc/core/CMakeLists.txt b/liger_cute_kernels/csrc/core/CMakeLists.txt index c326a030..03187cb0 100644 --- a/liger_cute_kernels/csrc/core/CMakeLists.txt +++ b/liger_cute_kernels/csrc/core/CMakeLists.txt @@ -20,15 +20,118 @@ list(FILTER LIGER_CUTE_CORE_SOURCES EXCLUDE REGEX "/src/moe/tune/") # The legacy TensorView ABI is no longer part of libliger_cute_kernels.so. # Tensor-carrying entry points are exported through TVM FFI. list(FILTER LIGER_CUTE_CORE_SOURCES EXCLUDE REGEX "/src/tensor_view\\.cpp$") + +# WGMMA/TMA translation units stay non-RDC. Remote NVSHMEM collectives are +# isolated in detail/remote_all_reduce.cu. +set(LIGER_CUTE_FSLCE_FORWARD_SM90_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu") +set(LIGER_CUTE_FSLCE_FORWARD_REDUCE_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy/forward_reduce.cu") +set(LIGER_CUTE_FSLCE_BACKWARD_SM90_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu") +set(LIGER_CUTE_FSLCE_BACKWARD_REDUCE_SM90_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu") +list(REMOVE_ITEM LIGER_CUTE_CORE_SOURCES + "${LIGER_CUTE_FSLCE_FORWARD_SM90_SOURCE}" + "${LIGER_CUTE_FSLCE_FORWARD_REDUCE_SOURCE}" + "${LIGER_CUTE_FSLCE_BACKWARD_SM90_SOURCE}" + "${LIGER_CUTE_FSLCE_BACKWARD_REDUCE_SM90_SOURCE}") + file(GLOB_RECURSE LIGER_CUTE_CORE_HEADERS "${CMAKE_CURRENT_SOURCE_DIR}/include/*.h" "${CMAKE_CURRENT_SOURCE_DIR}/include/*.cuh" ) list(FILTER LIGER_CUTE_CORE_HEADERS EXCLUDE REGEX "/include/liger_cute/tensor_view\\.h$") +add_library(liger_cute_fslce_forward_sm90 OBJECT + "${LIGER_CUTE_FSLCE_FORWARD_SM90_SOURCE}") +set_target_properties(liger_cute_fslce_forward_sm90 PROPERTIES + CUDA_SEPARABLE_COMPILATION OFF + CXX_VISIBILITY_PRESET hidden + CUDA_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) +target_compile_definitions(liger_cute_fslce_forward_sm90 PRIVATE + LIGER_CUTE_BUILDING=1) +target_include_directories(liger_cute_fslce_forward_sm90 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src/moe + ${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy + ${NCCL_INCLUDE_DIR}) +target_link_libraries(liger_cute_fslce_forward_sm90 PRIVATE + CUTLASS::CUTLASS + CUDA::cudart + ${NCCL_LIBRARY}) + +add_library(liger_cute_fslce_forward_reduce OBJECT + "${LIGER_CUTE_FSLCE_FORWARD_REDUCE_SOURCE}") +set_target_properties(liger_cute_fslce_forward_reduce PROPERTIES + CUDA_SEPARABLE_COMPILATION ON + CXX_VISIBILITY_PRESET hidden + CUDA_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) +target_compile_definitions(liger_cute_fslce_forward_reduce PRIVATE + LIGER_CUTE_BUILDING=1) +target_compile_options(liger_cute_fslce_forward_reduce PRIVATE + $<$:-UNVSHMEM_ENABLE_ALL_DEVICE_INLINING>) +target_include_directories(liger_cute_fslce_forward_reduce PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src/moe + ${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy) +target_link_libraries(liger_cute_fslce_forward_reduce PRIVATE + CUTLASS::CUTLASS + NVSHMEM::nvshmem_host + NVSHMEM::nvshmem_device + CUDA::cudart + CUDA::cuda_driver) + +add_library(liger_cute_fslce_backward_sm90 OBJECT + "${LIGER_CUTE_FSLCE_BACKWARD_SM90_SOURCE}") +set_target_properties(liger_cute_fslce_backward_sm90 PROPERTIES + CUDA_SEPARABLE_COMPILATION OFF + CXX_VISIBILITY_PRESET hidden + CUDA_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) +target_compile_definitions(liger_cute_fslce_backward_sm90 PRIVATE + LIGER_CUTE_BUILDING=1) +target_include_directories(liger_cute_fslce_backward_sm90 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src/moe + ${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy) +target_link_libraries(liger_cute_fslce_backward_sm90 PRIVATE + CUTLASS::CUTLASS + NVSHMEM::nvshmem_host + CUDA::cudart + CUDA::cuda_driver) + +add_library(liger_cute_fslce_backward_reduce_sm90 OBJECT + "${LIGER_CUTE_FSLCE_BACKWARD_REDUCE_SM90_SOURCE}") +set_target_properties(liger_cute_fslce_backward_reduce_sm90 PROPERTIES + CUDA_SEPARABLE_COMPILATION OFF + CXX_VISIBILITY_PRESET hidden + CUDA_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) +target_compile_definitions( + liger_cute_fslce_backward_reduce_sm90 PRIVATE + LIGER_CUTE_BUILDING=1) +target_include_directories( + liger_cute_fslce_backward_reduce_sm90 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src/moe + ${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy) +target_link_libraries( + liger_cute_fslce_backward_reduce_sm90 PRIVATE + CUTLASS::CUTLASS + NVSHMEM::nvshmem_host + CUDA::cudart + CUDA::cuda_driver) + add_library(liger_cute_kernels SHARED ${LIGER_CUTE_CORE_SOURCES} ${LIGER_CUTE_CORE_HEADERS} + $ + $ + $ + $ "${CMAKE_CURRENT_SOURCE_DIR}/../../liger_cute_kernels/tvm_ffi_bindings.cpp" ) @@ -45,13 +148,14 @@ set_target_properties(liger_cute_kernels PROPERTIES target_compile_definitions(liger_cute_kernels PRIVATE LIGER_CUTE_BUILDING=1) target_compile_options(liger_cute_kernels PRIVATE ${TVM_FFI_CFLAGS}) - target_include_directories(liger_cute_kernels PUBLIC $ $ PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src/moe + ${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy + ${NCCL_INCLUDE_DIR} ) # Deliberately NO torch here — that is the entire point of the split. @@ -64,6 +168,7 @@ target_link_libraries(liger_cute_kernels PUBLIC NVSHMEM::nvshmem_device CUDA::cudart CUDA::cuda_driver + ${NCCL_LIBRARY} ${TVM_FFI_LDFLAGS} tvm_ffi ) diff --git a/liger_cute_kernels/csrc/core/include/liger_cute/detail/local_reduce.cuh b/liger_cute_kernels/csrc/core/include/liger_cute/detail/local_reduce.cuh new file mode 100644 index 00000000..a319ad23 --- /dev/null +++ b/liger_cute_kernels/csrc/core/include/liger_cute/detail/local_reduce.cuh @@ -0,0 +1,154 @@ +#pragma once + +// Device-side local reduction algorithms for TpReducePlan backends. + +#include + +#include +#include +#include + +#include "liger_cute/detail/nvls.cuh" +#include "liger_cute/detail/tp_reduce.cuh" + +namespace liger_cute { +namespace detail { + +template +struct LocalReduceContext; + +template +struct LocalReduceContext { + std::uint64_t* local_ready; + std::uint64_t* multicast_ready; + std::uint64_t* local_complete; + std::uint64_t* multicast_complete; + int rank; + int size; +}; + +template +struct LocalReduceContext { + const Element* const* peer_sources; + std::size_t source_offset; + std::uint64_t* local_ready; + std::uint64_t* const* peer_sync; + std::size_t ready_offset; + std::uint64_t* local_complete; + std::size_t complete_offset; + int rank; + int size; +}; + +__device__ __forceinline__ void direct_peer_barrier_warp( + std::uint64_t* local_signals, + std::uint64_t* const* peer_signals, + std::size_t signal_offset, + int rank, + int size, + std::uint64_t epoch) { + unsigned int lane = nvls_lane_id(); + for (int peer = static_cast(lane); peer < size; peer += 32) { + asm volatile( + "st.release.sys.global.u64 [%0], %1;" + : + : "l"(peer_signals[peer] + signal_offset + rank), "l"(epoch) + : "memory"); + } + __syncwarp(); + + for (int peer = static_cast(lane); peer < size; peer += 32) { + std::uint64_t observed = 0; + do { + asm volatile( + "ld.acquire.sys.global.u64 %0, [%1];" + : "=l"(observed) + : "l"(local_signals + peer) + : "memory"); + } while (observed != epoch); + } + __syncwarp(); +} + +template +__device__ __forceinline__ void local_all_reduce( + const LocalReduceContext& context, + Element* destination, + const Element* source, + std::size_t count, + std::uint64_t epoch) { + if constexpr (Backend == LocalReduceBackend::kNvls) { + static_assert( + std::is_same_v, + "NVLS two-shot currently supports FP32 SUM"); + nvls_barrier_warp( + context.local_ready, + context.multicast_ready, + context.rank, + context.size, + epoch); + nvls_sum_reduce_warp_twoshot( + destination, + source, + count, + context.rank, + context.size); + nvls_barrier_warp( + context.local_complete, + context.multicast_complete, + context.rank, + context.size, + epoch); + } else { + static_assert( + std::is_same_v, + "direct-peer all-reduce currently supports FP32 SUM"); + direct_peer_barrier_warp( + context.local_ready, + context.peer_sync, + context.ready_offset, + context.rank, + context.size, + epoch); + + unsigned int lane = nvls_lane_id(); + std::size_t vectors = count / 4; + for (std::size_t vector = lane; vector < vectors; vector += 32) { + float4 sum = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int peer = 0; peer < context.size; ++peer) { + const auto* peer_source = + reinterpret_cast( + context.peer_sources[peer] + + context.source_offset); + float4 value = peer_source[vector]; + sum.x += value.x; + sum.y += value.y; + sum.z += value.z; + sum.w += value.w; + } + reinterpret_cast(destination)[vector] = sum; + } + for (std::size_t index = vectors * 4 + lane; + index < count; + index += 32) { + float sum = 0.0f; + for (int peer = 0; peer < context.size; ++peer) { + sum += context.peer_sources[peer][ + context.source_offset + index]; + } + destination[index] = sum; + } + __syncwarp(); + + direct_peer_barrier_warp( + context.local_complete, + context.peer_sync, + context.complete_offset, + context.rank, + context.size, + epoch); + } +} + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/include/liger_cute/detail/nvls.cuh b/liger_cute_kernels/csrc/core/include/liger_cute/detail/nvls.cuh new file mode 100644 index 00000000..5d1b84d7 --- /dev/null +++ b/liger_cute_kernels/csrc/core/include/liger_cute/detail/nvls.cuh @@ -0,0 +1,304 @@ +// Team-state-free NVLink SHARP primitives. +// +// Callers resolve the multicast pointers once, outside the consuming kernel, +// and pass them as ordinary kernel arguments. The helpers below therefore do +// not include NVSHMEM headers or reference NVSHMEM device state, allowing the +// consuming translation unit to remain non-RDC. +#pragma once + +#include + +#include +#include + +namespace liger_cute { +namespace detail { + +__device__ __forceinline__ unsigned int nvls_lane_id() { + unsigned int lane = 0; +#if defined(__CUDA_ARCH__) + asm("mov.u32 %0, %%laneid;" : "=r"(lane)); +#endif + return lane; +} + +// Barrier for one independent NVLS channel. +// +// Each PE multicast-stores its epoch into a distinct rank slot, then waits on +// the local replicas of every slot. Different CTAs may use the same multicast +// mapping concurrently as long as their signal vectors do not overlap. +__device__ __forceinline__ void nvls_barrier_warp( + std::uint64_t* local_signals, std::uint64_t* multicast_signals, + int team_rank, int team_size, std::uint64_t epoch) { +#if defined(__CUDA_ARCH__) + static_assert(__CUDA_ARCH__ >= 900, + "nvls_barrier_warp requires SM90 or newer"); + + const unsigned int lane = nvls_lane_id(); + if (lane == 0) { + asm volatile("multimem.st.release.sys.global.b64 [%0], %1;" + : + : "l"(multicast_signals + team_rank), "l"(epoch) + : "memory"); + } + __syncwarp(); + + for (int rank = static_cast(lane); rank < team_size; rank += 32) { + std::uint64_t observed = 0; + do { + asm volatile("ld.acquire.sys.global.u64 %0, [%1];" + : "=l"(observed) + : "l"(local_signals + rank) + : "memory"); + } while (observed != epoch); + } + __syncwarp(); +#else + (void)local_signals; + (void)multicast_signals; + (void)team_rank; + (void)team_size; + (void)epoch; +#endif +} + +// FP32 SUM all-reduce using the NVLS two-shot algorithm. +// +// Each PE reduces one disjoint slice of the region through the multicast +// source pointer, then multicast-stores that slice to every PE's destination. +// The combination of all corresponding warps therefore produces the complete +// result on every PE. +// +// This primitive intentionally owns no collective state. Concurrent CTAs may +// call it for disjoint, identically prefixed regions through the same TP team's +// multicast mapping. The caller must separately guarantee: +// +// * every PE has published its source region before any PE reads it; +// * one converged full warp per PE calls with warp-uniform arguments; +// * source and destination are not reused until every PE has completed; +// * multicast_source and multicast_dest belong to the same NVLS-capable team. +// * team_size is positive, team_rank is in [0, team_size), and both match the +// multicast mapping's membership; +// * count, rank assignment, and the region prefix are identical across PEs; +// * nonempty regions are at least four-byte aligned and the source/destination +// regions are either exactly in-place or non-overlapping. +// +// The pointers passed here must already include any block/channel prefix. +// A cross-PE completion protocol is still required after this function returns. +// Requires an NVLS-capable SM90-or-newer target. +__device__ __forceinline__ void nvls_sum_reduce_warp_twoshot( + float* multicast_dest, const float* multicast_source, std::size_t count, + int team_rank, int team_size) { +#if defined(__CUDA_ARCH__) + static_assert(__CUDA_ARCH__ >= 900, + "nvls_sum_reduce_warp_twoshot requires SM90 or newer"); + + if (count == 0) return; + + const unsigned int lane = nvls_lane_id(); + const std::size_t elements_per_pe = + count / static_cast(team_size); + const std::size_t rank_offset = + elements_per_pe * static_cast(team_rank); + const std::size_t local_count = + elements_per_pe + + (team_rank == team_size - 1 + ? count % static_cast(team_size) + : 0); + + const float* source = multicast_source + rank_offset; + float* dest = multicast_dest + rank_offset; + std::size_t remaining = local_count; + + __syncwarp(); + asm volatile("" ::: "memory"); + + if ((reinterpret_cast(source) & 0xF) == 0 && + (reinterpret_cast(dest) & 0xF) == 0) { + const std::size_t vectors = remaining / 4; + const auto* source_vectors = reinterpret_cast(source); + auto* dest_vectors = reinterpret_cast(dest); + for (std::size_t index = static_cast(lane); index < vectors; + index += 32) { + std::uint32_t value[4]; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 " + "{%0, %1, %2, %3}, [%4];" + : "=r"(value[0]), "=r"(value[1]), "=r"(value[2]), "=r"(value[3]) + : "l"(source_vectors + index) + : "memory"); + asm volatile( + "multimem.st.relaxed.sys.global.v4.f32 " + "[%0], {%1, %2, %3, %4};" + : + : "l"(dest_vectors + index), "r"(value[0]), "r"(value[1]), + "r"(value[2]), "r"(value[3]) + : "memory"); + } + const std::size_t consumed = vectors * 4; + source += consumed; + dest += consumed; + remaining -= consumed; + } + + if ((reinterpret_cast(source) & 0x7) == 0 && + (reinterpret_cast(dest) & 0x7) == 0) { + const std::size_t vectors = remaining / 2; + const auto* source_vectors = + reinterpret_cast(source); + auto* dest_vectors = reinterpret_cast(dest); + for (std::size_t index = static_cast(lane); index < vectors; + index += 32) { + std::uint32_t value[2]; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.v2.f32 " + "{%0, %1}, [%2];" + : "=r"(value[0]), "=r"(value[1]) + : "l"(source_vectors + index) + : "memory"); + asm volatile( + "multimem.st.relaxed.sys.global.v2.f32 [%0], {%1, %2};" + : + : "l"(dest_vectors + index), "r"(value[0]), "r"(value[1]) + : "memory"); + } + const std::size_t consumed = vectors * 2; + source += consumed; + dest += consumed; + remaining -= consumed; + } + + for (std::size_t index = static_cast(lane); index < remaining; + index += 32) { + std::uint32_t value; + asm volatile("multimem.ld_reduce.relaxed.sys.global.add.f32 %0, [%1];" + : "=r"(value) + : "l"(source + index) + : "memory"); + asm volatile("multimem.st.relaxed.sys.global.f32 [%0], %1;" + : + : "l"(dest + index), "r"(value) + : "memory"); + } + + asm volatile("fence.acq_rel.sys;" ::: "memory"); + __syncwarp(); +#else + (void)multicast_dest; + (void)multicast_source; + (void)count; + (void)team_rank; + (void)team_size; +#endif +} + +__device__ __forceinline__ std::size_t nvls_rank_count( + std::size_t count, int rank, int size) { + std::size_t per_rank = count / static_cast(size); + return per_rank + + (rank == size - 1 ? count % static_cast(size) : 0); +} + +__device__ __forceinline__ std::size_t nvls_rank_offset( + std::size_t count, int rank, int size) { + return count / static_cast(size) * + static_cast(rank); +} + +__device__ __forceinline__ void nvls_sum_reduce_scatter_warp( + float* local_dest, const float* multicast_source, std::size_t count, + int team_rank, int team_size) { +#if defined(__CUDA_ARCH__) + const unsigned int lane = nvls_lane_id(); + const std::size_t rank_offset = + nvls_rank_offset(count, team_rank, team_size); + const std::size_t local_count = + nvls_rank_count(count, team_rank, team_size); + const float* source = multicast_source + rank_offset; + + const std::size_t vectors = local_count / 4; + for (std::size_t index = lane; index < vectors; index += 32) { + std::uint32_t value[4]; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 " + "{%0, %1, %2, %3}, [%4];" + : "=r"(value[0]), "=r"(value[1]), "=r"(value[2]), "=r"(value[3]) + : "l"(reinterpret_cast(source) + index) + : "memory"); + asm volatile( + "st.global.v4.b32 [%0], {%1, %2, %3, %4};" + : + : "l"(reinterpret_cast(local_dest) + index), "r"(value[0]), + "r"(value[1]), "r"(value[2]), "r"(value[3]) + : "memory"); + } + for (std::size_t index = vectors * 4 + lane; index < local_count; + index += 32) { + std::uint32_t value; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.f32 %0, [%1];" + : "=r"(value) + : "l"(source + index) + : "memory"); + asm volatile( + "st.global.b32 [%0], %1;" + : + : "l"(local_dest + index), "r"(value) + : "memory"); + } + asm volatile("fence.acq_rel.sys;" ::: "memory"); + __syncwarp(); +#else + (void)local_dest; + (void)multicast_source; + (void)count; + (void)team_rank; + (void)team_size; +#endif +} + +__device__ __forceinline__ void nvls_allgather_warp( + float* multicast_dest, const float* local_source, std::size_t count, + int team_rank, int team_size) { +#if defined(__CUDA_ARCH__) + const unsigned int lane = nvls_lane_id(); + const std::size_t rank_offset = + nvls_rank_offset(count, team_rank, team_size); + const std::size_t local_count = + nvls_rank_count(count, team_rank, team_size); + float* destination = multicast_dest + rank_offset; + + const std::size_t vectors = local_count / 4; + for (std::size_t index = lane; index < vectors; index += 32) { + int4 value = reinterpret_cast(local_source)[index]; + asm volatile( + "multimem.st.relaxed.sys.global.v4.f32 " + "[%0], {%1, %2, %3, %4};" + : + : "l"(reinterpret_cast(destination) + index), "r"(value.x), + "r"(value.y), "r"(value.z), "r"(value.w) + : "memory"); + } + for (std::size_t index = vectors * 4 + lane; index < local_count; + index += 32) { + std::uint32_t value = + reinterpret_cast(local_source)[index]; + asm volatile( + "multimem.st.relaxed.sys.global.f32 [%0], %1;" + : + : "l"(destination + index), "r"(value) + : "memory"); + } + asm volatile("fence.acq_rel.sys;" ::: "memory"); + __syncwarp(); +#else + (void)multicast_dest; + (void)local_source; + (void)count; + (void)team_rank; + (void)team_size; +#endif +} + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/include/liger_cute/detail/tp_reduce.cuh b/liger_cute_kernels/csrc/core/include/liger_cute/detail/tp_reduce.cuh new file mode 100644 index 00000000..d70a3389 --- /dev/null +++ b/liger_cute_kernels/csrc/core/include/liger_cute/detail/tp_reduce.cuh @@ -0,0 +1,111 @@ +#pragma once + +#include + +#include +#include + +namespace liger_cute { +namespace detail { + +inline constexpr int kMaxTpReduceTeamSize = 512; + +enum class LocalReduceBackend : std::uint8_t { + kNvls, + kDirectPeer, +}; + +// Minimal device view for the local NVLS reduce-scatter/all-gather path. +struct NvlsReduceView { + float* multicast_partial; + float* multicast_reduced; + std::uint64_t* multicast_sync; + float* reduced_shard; + int rank; + int size; +}; + +// Minimal device view for the non-NVLS direct-peer all-reduce path. +struct DirectPeerReduceView { + const float* const* peer_partial; + std::uint64_t* const* peer_sync; + int available; + int rank; + int size; +}; + +// Host-only description of the optional follow-up remote reduction. +struct RemoteReduceView { + float* reduced_shard; + float* inbox; + std::uint64_t* ready; + std::uint64_t* consumed; + int rank; + int size; + int peer_world; + + bool enabled() const { + return peer_world >= 0; + } +}; + +// Host dispatch plan. Only the selected backend-specific view is passed to a +// kernel; RequiresRemote remains a compile-time launcher parameter. +struct TpReducePlan { + LocalReduceBackend backend; + NvlsReduceView nvls; + DirectPeerReduceView direct; + RemoteReduceView remote; + int team_size; +}; + +struct TpReduceTopology { + int team_size; + int local_size; +}; + +struct TpReduceBuffers { + float* partial; + float* reduced; + float* reduced_shard; + float* remote_inbox; + std::uint64_t* remote_signals; + std::uint64_t* sync; + std::size_t sync_bytes; + float** peer_partial_storage; + std::uint64_t** peer_sync_storage; +}; + +#if defined(__CUDACC__) +__device__ __forceinline__ void publish_local_reduce_source() { + asm volatile("fence.proxy.async.global;" ::: "memory"); + __threadfence_system(); +} +#endif + +// Collective across parent_team. Resolves NVLS and direct-peer mappings once +// and selects the local backend. A two-host remote stage is represented +// separately in TpReducePlan::remote. +void configure_tp_reduce( + std::int64_t parent_team, + const TpReduceBuffers& buffers); + +TpReduceTopology query_tp_reduce_topology(std::int64_t parent_team); +TpReducePlan tp_reduce_plan(); + +void begin_tp_reduce( + const std::uint64_t* launch_epoch, + cudaStream_t stream); +void end_tp_reduce(cudaStream_t stream); +void reset_tp_reduce(); + +// Separate RDC follow-up stage for the packed shard produced by the local +// reduction. This is intentionally outside the hot non-RDC WGMMA kernel. +void launch_remote_reduce( + const RemoteReduceView& remote, + const std::uint64_t* launch_epoch, + std::size_t count, + cudaStream_t stream); + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cu b/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cu new file mode 100644 index 00000000..96581174 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cu @@ -0,0 +1,84 @@ +#include "remote_all_reduce.cuh" + +#include +#include + +#include "liger_cute/check.h" + +namespace liger_cute { +namespace detail { +namespace { + +__global__ void remote_put( + float* inbox, + const float* source, + std::size_t count, + int peer) { + nvshmemx_float_put_block(inbox, source, count, peer); +} + +__global__ void remote_handshake( + std::uint64_t* signal, + const std::uint64_t* launch_epoch, + std::uint64_t epoch_suffix, + int peer) { + std::uint64_t epoch = *launch_epoch | epoch_suffix; + nvshmemx_signal_op( + signal, epoch, NVSHMEM_SIGNAL_SET, peer); + nvshmem_quiet(); + nvshmem_signal_wait_until(signal, NVSHMEM_CMP_EQ, epoch); +} + +__global__ void remote_sum( + float* destination, + const float* source, + const float* inbox, + std::size_t count) { + for (std::size_t index = + blockIdx.x * blockDim.x + threadIdx.x; + index < count; + index += blockDim.x * gridDim.x) { + destination[index] = source[index] + inbox[index]; + } +} + +} // namespace + +void launch_remote_pair_all_reduce( + float* destination, + float* inbox, + const float* source, + std::size_t count, + std::uint64_t* ready, + std::uint64_t* consumed, + const std::uint64_t* launch_epoch, + std::uint64_t epoch_suffix, + int peer, + cudaStream_t stream) { + LIGER_CHECK( + destination != nullptr && inbox != nullptr && source != nullptr && + ready != nullptr && consumed != nullptr && + launch_epoch != nullptr, + "remote all-reduce buffers must be non-null"); + LIGER_CHECK(count > 0, "remote all-reduce count must be positive"); + LIGER_CHECK(peer >= 0, "remote all-reduce peer must be non-negative"); + + remote_put<<<1, 256, 0, stream>>>(inbox, source, count, peer); + nvshmemx_quiet_on_stream(stream); + remote_handshake<<<1, 1, 0, stream>>>( + ready, launch_epoch, epoch_suffix, peer); + int blocks = static_cast( + (count + 255) / 256 < 1024 ? (count + 255) / 256 : 1024); + remote_sum<<>>( + destination, source, inbox, count); + remote_handshake<<<1, 1, 0, stream>>>( + consumed, launch_epoch, epoch_suffix, peer); + cudaError_t error = cudaGetLastError(); + LIGER_CHECK( + error == cudaSuccess, + "remote pair all-reduce launch failed: ", + cudaGetErrorString(error)); +} + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cuh b/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cuh new file mode 100644 index 00000000..24240a07 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cuh @@ -0,0 +1,26 @@ +#pragma once + +// Internal implementation detail behind launch_remote_reduce(). + +#include + +#include +#include + +namespace liger_cute { +namespace detail { + +void launch_remote_pair_all_reduce( + float* destination, + float* inbox, + const float* source, + std::size_t count, + std::uint64_t* ready, + std::uint64_t* consumed, + const std::uint64_t* launch_epoch, + std::uint64_t epoch_suffix, + int peer, + cudaStream_t stream); + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/src/detail/tp_reduce.cu b/liger_cute_kernels/csrc/core/src/detail/tp_reduce.cu new file mode 100644 index 00000000..5ae1ee7b --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/detail/tp_reduce.cu @@ -0,0 +1,357 @@ +#include "liger_cute/detail/tp_reduce.cuh" + +#include +#include +#include + +#include +#include +#include + +#include "remote_all_reduce.cuh" +#include "liger_cute/check.h" + +namespace liger_cute { +namespace detail { +namespace { + +struct RawNvlsMapping { + float* multicast_partial; + float* multicast_reduced; + std::uint64_t* multicast_sync; + float* node_multicast_partial; + float* node_multicast_reduced; + std::uint64_t* node_multicast_sync; + int team_rank; + int team_size; + int node_rank; + int node_size; + int remote_rank; + int remote_size; +}; + +RawNvlsMapping g_raw = {}; +TpReducePlan g_plan = {}; +TpReduceBuffers g_buffers = {}; +bool g_configured = false; +std::int64_t g_parent_team = 0; + +void check_cuda(cudaError_t error, const char* what) { + LIGER_CHECK( + error == cudaSuccess, + "tensor-parallel reduction: ", + what, + " failed: ", + cudaGetErrorString(error)); +} + +__global__ void query_nvls_mapping( + nvshmem_team_t team, + float* partial, + float* reduced, + std::uint64_t* sync, + RawNvlsMapping* output) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + output->multicast_partial = + static_cast(nvshmemx_mc_ptr(team, partial)); + output->multicast_reduced = + static_cast(nvshmemx_mc_ptr(team, reduced)); + output->multicast_sync = + static_cast(nvshmemx_mc_ptr(team, sync)); + output->node_multicast_partial = + static_cast( + nvshmemx_mc_ptr(NVSHMEMX_TEAM_NODE, partial)); + output->node_multicast_reduced = + static_cast( + nvshmemx_mc_ptr(NVSHMEMX_TEAM_NODE, reduced)); + output->node_multicast_sync = + static_cast( + nvshmemx_mc_ptr(NVSHMEMX_TEAM_NODE, sync)); + output->team_rank = nvshmem_team_my_pe(team); + output->team_size = nvshmem_team_n_pes(team); + output->node_rank = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + output->node_size = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + output->remote_rank = + nvshmem_team_my_pe(NVSHMEMX_TEAM_SAME_MYPE_NODE); + output->remote_size = + nvshmem_team_n_pes(NVSHMEMX_TEAM_SAME_MYPE_NODE); +} + +__global__ void advance_launch_epoch(std::uint64_t* launch_epoch) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + *launch_epoch += std::uint64_t{1} << 32; + } +} + +bool same_buffers(const TpReduceBuffers& lhs, const TpReduceBuffers& rhs) { + return lhs.partial == rhs.partial && + lhs.reduced == rhs.reduced && + lhs.reduced_shard == rhs.reduced_shard && + lhs.remote_inbox == rhs.remote_inbox && + lhs.remote_signals == rhs.remote_signals && + lhs.sync == rhs.sync && + lhs.sync_bytes == rhs.sync_bytes && + lhs.peer_partial_storage == rhs.peer_partial_storage && + lhs.peer_sync_storage == rhs.peer_sync_storage; +} + +} // namespace + +TpReduceTopology query_tp_reduce_topology(std::int64_t parent_team) { + nvshmem_team_t team = static_cast(parent_team); + int team_size = nvshmem_team_n_pes(team); + int node_size = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + LIGER_CHECK( + team_size >= 1 && team_size <= kMaxTpReduceTeamSize, + "invalid tensor-parallel reduction team size ", + team_size); + LIGER_CHECK(node_size >= 1, "invalid NVSHMEM node-team size"); + return { + team_size, + team_size < node_size ? team_size : node_size, + }; +} + +void configure_tp_reduce( + std::int64_t parent_team, + const TpReduceBuffers& buffers) { + LIGER_CHECK( + buffers.partial != nullptr && buffers.reduced != nullptr && + buffers.reduced_shard != nullptr && + buffers.remote_inbox != nullptr && + buffers.remote_signals != nullptr && + buffers.sync != nullptr, + "tensor-parallel reduction buffers must be non-null"); + LIGER_CHECK( + buffers.sync_bytes > 0, + "tensor-parallel reduction signal storage must be non-empty"); + LIGER_CHECK( + buffers.peer_partial_storage != nullptr && + buffers.peer_sync_storage != nullptr, + "tensor-parallel direct-peer mapping storage must be non-null"); + + if (g_configured) { + LIGER_CHECK( + parent_team == g_parent_team && + same_buffers(buffers, g_buffers), + "tensor-parallel reduction mapping is immutable once configured"); + return; + } + + nvshmem_team_t team = static_cast(parent_team); + int team_size = query_tp_reduce_topology(parent_team).team_size; + + check_cuda( + cudaMemset(buffers.sync, 0, buffers.sync_bytes), + "cudaMemset(reduction signals)"); + check_cuda( + cudaMemset( + buffers.remote_signals, + 0, + 2 * sizeof(std::uint64_t)), + "cudaMemset(remote reduction signals)"); + + RawNvlsMapping* device_mapping = nullptr; + check_cuda( + cudaMalloc(&device_mapping, sizeof(RawNvlsMapping)), + "cudaMalloc(reduction mapping)"); + query_nvls_mapping<<<1, 1>>>( + team, + buffers.partial, + buffers.reduced, + buffers.sync, + device_mapping); + check_cuda(cudaGetLastError(), "query_nvls_mapping launch"); + check_cuda( + cudaMemcpy( + &g_raw, + device_mapping, + sizeof(RawNvlsMapping), + cudaMemcpyDeviceToHost), + "cudaMemcpy(reduction mapping)"); + check_cuda(cudaFree(device_mapping), "cudaFree(reduction mapping)"); + + LIGER_CHECK( + g_raw.team_rank >= 0 && + g_raw.team_rank < g_raw.team_size && + g_raw.team_size == team_size, + "inconsistent tensor-parallel reduction team metadata"); + + std::vector peer_partial(team_size); + std::vector peer_sync(team_size); + std::vector world_members( + static_cast(nvshmem_n_pes()), 0); + bool direct_available = true; + bool parent_covers_world = team_size == nvshmem_n_pes(); + int my_world_pe = nvshmem_my_pe(); + for (int rank = 0; rank < team_size; ++rank) { + int world_pe = nvshmem_team_translate_pe( + team, rank, NVSHMEM_TEAM_WORLD); + LIGER_CHECK( + world_pe >= 0, + "failed to translate tensor-parallel rank ", + rank, + " to NVSHMEM_TEAM_WORLD"); + if (world_pe >= static_cast(world_members.size()) || + world_members[world_pe] != 0) { + parent_covers_world = false; + } else { + world_members[world_pe] = 1; + } + peer_partial[rank] = world_pe == my_world_pe + ? buffers.partial + : static_cast( + nvshmem_ptr(buffers.partial, world_pe)); + peer_sync[rank] = world_pe == my_world_pe + ? buffers.sync + : static_cast( + nvshmem_ptr(buffers.sync, world_pe)); + direct_available = + direct_available && peer_partial[rank] != nullptr && + peer_sync[rank] != nullptr; + } + check_cuda( + cudaMemcpy( + buffers.peer_partial_storage, + peer_partial.data(), + peer_partial.size() * sizeof(float*), + cudaMemcpyHostToDevice), + "cudaMemcpy(direct-peer partial pointers)"); + check_cuda( + cudaMemcpy( + buffers.peer_sync_storage, + peer_sync.data(), + peer_sync.size() * sizeof(std::uint64_t*), + cudaMemcpyHostToDevice), + "cudaMemcpy(direct-peer signal pointers)"); + bool nvls_available = + team_size == 1 || + (g_raw.multicast_partial != nullptr && + g_raw.multicast_reduced != nullptr && + g_raw.multicast_sync != nullptr); + bool remote_available = + !nvls_available && parent_covers_world && + g_raw.node_size > 1 && + g_raw.remote_size == 2 && + g_raw.node_size * g_raw.remote_size == team_size && + g_raw.node_multicast_partial != nullptr && + g_raw.node_multicast_reduced != nullptr && + g_raw.node_multicast_sync != nullptr; + + g_plan = {}; + g_plan.remote.peer_world = -1; + g_plan.team_size = g_raw.team_size; + g_plan.direct = { + buffers.peer_partial_storage, + buffers.peer_sync_storage, + static_cast(direct_available), + g_raw.team_rank, + g_raw.team_size}; + + if (nvls_available) { + g_plan.backend = LocalReduceBackend::kNvls; + g_plan.nvls = { + team_size == 1 ? buffers.partial : g_raw.multicast_partial, + team_size == 1 ? buffers.reduced : g_raw.multicast_reduced, + team_size == 1 ? buffers.sync : g_raw.multicast_sync, + buffers.reduced_shard, + g_raw.team_rank, + g_raw.team_size}; + } else if (remote_available) { + g_plan.backend = LocalReduceBackend::kNvls; + g_plan.nvls = { + g_raw.node_multicast_partial, + g_raw.node_multicast_reduced, + g_raw.node_multicast_sync, + buffers.reduced_shard, + g_raw.node_rank, + g_raw.node_size}; + int peer_rank = 1 - g_raw.remote_rank; + int peer_world = nvshmem_team_translate_pe( + NVSHMEMX_TEAM_SAME_MYPE_NODE, + peer_rank, + NVSHMEM_TEAM_WORLD); + LIGER_CHECK( + peer_world >= 0, + "failed to translate remote reduction peer to world rank"); + g_plan.remote = { + buffers.reduced_shard, + buffers.remote_inbox, + buffers.remote_signals, + buffers.remote_signals + 1, + g_raw.remote_rank, + g_raw.remote_size, + peer_world}; + } else { + g_plan.backend = LocalReduceBackend::kDirectPeer; + } + + int barrier_status = nvshmem_barrier(team); + LIGER_CHECK( + barrier_status == 0, + "tensor-parallel reduction setup barrier failed with status ", + barrier_status); + + g_parent_team = parent_team; + g_buffers = buffers; + g_configured = true; +} + +TpReducePlan tp_reduce_plan() { + LIGER_CHECK( + g_configured, + "configure_tp_reduce() must be called before requesting a plan"); + return g_plan; +} + +void begin_tp_reduce( + const std::uint64_t* launch_epoch, + cudaStream_t stream) { + LIGER_CHECK(g_configured, "tensor-parallel reduction is not configured"); + LIGER_CHECK(launch_epoch != nullptr, "launch epoch is null"); + advance_launch_epoch<<<1, 1, 0, stream>>>( + const_cast(launch_epoch)); + check_cuda(cudaGetLastError(), "advance_launch_epoch launch"); +} + +void end_tp_reduce(cudaStream_t stream) { + LIGER_CHECK(g_configured, "tensor-parallel reduction is not configured"); + if (g_plan.team_size > 1) { + nvshmemx_barrier_on_stream( + static_cast(g_parent_team), stream); + } +} + +void launch_remote_reduce( + const RemoteReduceView& remote, + const std::uint64_t* launch_epoch, + std::size_t count, + cudaStream_t stream) { + LIGER_CHECK(remote.enabled(), "remote reduction is not configured"); + LIGER_CHECK( + remote.size == 2 && remote.rank >= 0 && remote.rank < remote.size, + "invalid remote reduction topology"); + launch_remote_pair_all_reduce( + remote.reduced_shard, + remote.inbox, + remote.reduced_shard, + count, + remote.ready, + remote.consumed, + launch_epoch, + 0x30000000u, + remote.peer_world, + stream); +} + +void reset_tp_reduce() { + g_raw = {}; + g_plan = {}; + g_plan.remote.peer_world = -1; + g_buffers = {}; + g_configured = false; + g_parent_team = 0; +} + +} // namespace detail +} // namespace liger_cute diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward.cuh new file mode 100644 index 00000000..c451206c --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward.cuh @@ -0,0 +1,47 @@ +#pragma once + +#include + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +__host__ __device__ inline float scaled_cross_entropy_dlogit( + float scaled_logit, + float lse, + bool is_target, + float nll_scale, + float temperature, + float expected_scaled_logit = 0.0f, + float entropy_scale = 0.0f, + bool ignored = false) { + if (ignored) return 0.0f; + + float probability = expf(scaled_logit - lse); + float gradient = nll_scale * (probability - (is_target ? 1.0f : 0.0f)); + if (entropy_scale != 0.0f) { + gradient += entropy_scale * probability * + (expected_scaled_logit - scaled_logit); + } + return gradient / temperature; +} + +struct BackwardWave { + int first_m_tile; + int num_m_tiles; +}; + +template +__host__ __device__ constexpr BackwardWave make_backward_wave( + int wave, int total_m_tiles) { + int first_m_tile = wave * GemmConfig::kMTilesPerWave; + int remaining = total_m_tiles - first_m_tile; + int num_m_tiles = remaining < GemmConfig::kMTilesPerWave + ? remaining + : GemmConfig::kMTilesPerWave; + return {first_m_tile, num_m_tiles > 0 ? num_m_tiles : 0}; +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm90.cuh new file mode 100644 index 00000000..82884604 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm90.cuh @@ -0,0 +1,1907 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy backward kernels. +// +// dZ uses the cluster-2 M-fast handoff kernel: warp 3 produces TMA panels, +// warps 4..11 run N256 WGMMA, and warps 0..2 transform/store the staged logits. +// dX+dW uses one cluster-2 four-stage mainloop and one R24 accumulator reused +// across phases. dX publishes FP32 N32 TMA tiles to the communication ring; +// dW uses N64/S2 TMA epilogues with hidden-N pairs as the fastest raster. +// ═══════════════════════════════════════════════════════════════════════════ + +// gemm_sm90.cuh carries the reusable SM90 producer/consumer mechanics and +// enables CuTe's extended MMA shapes, so it must be the first CUTLASS/CuTe +// include of any TU using this header. +#include "gemm_sm90.cuh" + +#include "backward.cuh" +#include "backward_gemm_sm90.cuh" +#include "dx_reduce.cuh" +#include "tma_store_sm90.cuh" +#include "liger_cute/detail/local_reduce.cuh" + +#include +#include +#include + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +using namespace cute; + +// ─────────────────────────────────────────────────────────────────────────── +// SMEM layouts that gemm_sm90.cuh does not carry +// ─────────────────────────────────────────────────────────────────────────── + +// MN-major operand staging with the 128 B swizzle atom. dX's B operand (W seen +// as [hidden, vocab]) and both dW operands are MN-major, which the K-major +// helper in gemm_sm90.cuh cannot express. +template +struct SmemLayoutMnMajorSm90 { + static_assert(Compute == 90, "SmemLayoutMnMajorSm90 requires Compute=90"); + static_assert(Rows * static_cast(sizeof(Element)) % 128 == 0, + "the SW128 MN-major atom needs 128 B columns"); + + static constexpr int kCompute = Compute; + + using Atom = GMMA::Layout_MN_SW128_Atom; + using Single = decltype(tile_to_shape( + Atom{}, Shape, Int>{})); + using Staged = decltype(tile_to_shape( + Atom{}, Shape, Int, Int>{})); + + static constexpr int kStageBytes = static_cast( + cute::size(Single{}) * sizeof(Element)); + static constexpr int kStagedElements = cosize_v; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Global tensors for the TMA descriptors +// ─────────────────────────────────────────────────────────────────────────── + +// int64 extents throughout: tokens * hidden and vocab * hidden both exceed +// INT_MAX at production shapes and would wrap CuTe's layout algebra. +template +CUTE_HOST_DEVICE auto backward_row_major_tensor( + const Element* ptr, int64_t rows, int64_t cols) { + return make_tensor( + make_gmem_ptr(ptr), + make_layout(make_shape(rows, cols), make_stride(cols, _1{}))); +} + +// A transposed *view* of a row-major [cols, rows] matrix: shape (rows, cols) +// with the first mode contiguous. This is how the MN-major operands are +// addressed without ever materialising a transpose. +template +CUTE_HOST_DEVICE auto backward_col_major_tensor( + const Element* ptr, int64_t rows, int64_t cols, int64_t leading) { + return make_tensor( + make_gmem_ptr(ptr), + make_layout(make_shape(rows, cols), make_stride(_1{}, leading))); +} + +// ─────────────────────────────────────────────────────────────────────────── +// Traits +// ─────────────────────────────────────────────────────────────────────────── + +template < + int Compute = 90, + int MainloopStages = BackwardGemmConfigSm90::kMainloopStages, + int DwStoreTileN = BackwardGemmConfigSm90::kDwStoreTileN, + int DwStoreStages = BackwardGemmConfigSm90::kDwStoreStages> +struct BackwardGemmTraitsSm90 { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + using Element = cutlass::bfloat16_t; + using ElementAccum = float; + + static constexpr int kTileM = Config::kTileM; + static constexpr int kTileK = Config::kTileK; + static constexpr int kTileN = Config::kDzTileN; + static constexpr int kStages = MainloopStages; + static constexpr int kNumThreads = Config::kNumThreads; + static constexpr int kWarpSize = fused_scaled_linear_cross_entropy::kWarpSize; + static constexpr int kWarpGroupSize = Config::kWarpGroupSize; + static constexpr int kNumMmaWarpGroups = Config::kNumMmaWarpGroups; + static constexpr int kConsumerThreads = Config::kConsumerThreads; + static constexpr int kDwStoreTileN = DwStoreTileN; + static constexpr int kDwStoreStages = DwStoreStages; + + static_assert(Config::kDzTileN == Config::kDxTileN); + static_assert(Config::kDzTileN == Config::kDwTileN); + static_assert(Config::kClusterM == 1, + "nvshmemx_collective_launch() takes no cluster launch attribute, so " + "the backward kernel cannot run in a cluster"); + + using ClusterShape = Shape<_1, _1, _1>; + + // Operand staging. A is M128xK64, B is N256xK64; both are three stages + // deep and every phase re-uses the same two SMEM arenas. + using SmemAKMajor = + sm90::SmemLayoutKMajorSm90; + using SmemBKMajor = + sm90::SmemLayoutKMajorSm90; + using SmemAMnMajor = + SmemLayoutMnMajorSm90; + using SmemBMnMajor = + SmemLayoutMnMajorSm90; + + static constexpr int kSmemAElements = SmemAKMajor::kStagedElements; + static constexpr int kSmemBElements = SmemBKMajor::kStagedElements; + static_assert(kSmemAElements == SmemAMnMajor::kStagedElements, + "the K-major and MN-major A arenas must be interchangeable"); + static_assert(kSmemBElements == SmemBMnMajor::kStagedElements, + "the K-major and MN-major B arenas must be interchangeable"); + + // dX stores accumulator fragments directly to FP32 global staging; dW + // uses this shared-memory TMA epilogue. + using DwSmemStore1 = sm90::SmemLayoutKMajorSm90< + Element, kTileM, kDwStoreTileN, 1, Compute>; + template + using DwSmemStore = sm90::SmemLayoutKMajorSm90< + Element, kTileM, kDwStoreTileN, Stages, Compute>; + + static constexpr int kDwStoreElements = + DwSmemStore::kStagedElements; + static constexpr int kStoreElements = kDwStoreElements; + + // WGMMA. Two warp groups split the M128 tile; one N256 atom per group. + template + using TiledMmaTraits = sm90::TiledMmaMSplitSm90< + Element, Element, ElementAccum, kTileN, kNumMmaWarpGroups, + MajorA, MajorB, Compute>; + + using DxMmaTraits = TiledMmaTraits; + using DwMmaTraits = TiledMmaTraits; + using DwStoreMmaTraits = sm90::TiledMmaMSplitSm90< + Element, + Element, + ElementAccum, + kDwStoreTileN, + kNumMmaWarpGroups, + GMMA::Major::MN, + GMMA::Major::MN, + Compute>; + using AccumShape = Shape, Int>; + + using PipelineTraits = sm90::MainloopPipelineSm90; + using MainloopPipeline = typename PipelineTraits::Type; + using PipelineState = typename PipelineTraits::State; + + // One mbarrier per stage covers A and B together. + static constexpr int kTmaTransBytes = + SmemAKMajor::kStageBytes + SmemBKMajor::kStageBytes; + static_assert( + SmemAKMajor::kStageBytes == SmemAMnMajor::kStageBytes && + SmemBKMajor::kStageBytes == SmemBMnMajor::kStageBytes, + "all three phases must move the same number of bytes per stage"); + +}; + + +template +struct BackwardDxCluster2TraitsSm90 { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + using Config = BackwardGemmConfigSm90; + using Element = cutlass::bfloat16_t; + using ElementAccum = float; + + static constexpr int kTileM = 128; + static constexpr int kTileK = 64; + static constexpr int kTileN = 256; + static constexpr int kFragmentTileN = 256; + static constexpr int kNumFragments = 1; + static constexpr int kStoreTileN = 32; + static constexpr int kStoreStages = 1; + static constexpr int kStages = 4; + static constexpr int kClusterM = 2; + static constexpr int kNumThreads = 384; + static constexpr int kConsumerThreads = 256; + static constexpr int kWarpSize = 32; + + using ClusterShape = Shape, _1, _1>; + using SmemA = + sm90::SmemLayoutKMajorSm90< + Element, kTileM, kTileK, kStages, Compute>; + using SmemB = SmemLayoutMnMajorSm90< + Element, kFragmentTileN, kTileK, kStages, Compute>; + using SmemStore = sm90::SmemLayoutKMajorSm90< + float, kTileM, kStoreTileN, 1, Compute>; + using TiledMmaTraits = sm90::TiledMmaMSplitSm90< + Element, + Element, + ElementAccum, + kFragmentTileN, + 2, + GMMA::Major::K, + GMMA::Major::MN, + Compute>; + using TiledMma = typename TiledMmaTraits::Type; + static_assert(cute::is_same_v< + typename TiledMmaTraits::Atom, + SM90_64x256x16_F32BF16BF16_SS< + GMMA::Major::K, GMMA::Major::MN>>, + "production clustered dX must use native N256 WGMMA"); + using AccumShape = Shape, Int>; + using StoreMmaTraits = sm90::TiledMmaMSplitSm90< + Element, + Element, + float, + kStoreTileN, + 2, + GMMA::Major::K, + GMMA::Major::MN, + Compute>; + + using PipelineTraits = sm90::MainloopPipelineSm90; + using MainloopPipeline = typename PipelineTraits::Type; + using PipelineState = typename PipelineTraits::State; + + static constexpr int kSmemAElements = SmemA::kStagedElements; + static constexpr int kSmemBElements = SmemB::kStagedElements; + static constexpr int kTmaTransBytes = + SmemA::kStageBytes + kNumFragments * SmemB::kStageBytes; + + static_assert(kTileN == kNumFragments * kFragmentTileN); +}; + +template +struct BackwardDzHandoffTraitsSm90 { + using Config = BackwardGemmConfigSm90; + using Element = cutlass::bfloat16_t; + using ElementLogit = cutlass::half_t; + using ElementAccum = float; + + static constexpr int kTileM = 128; + static constexpr int kTileN = 256; + static constexpr int kTileK = 64; + static constexpr int kStages = 3; + static constexpr int kStoreTileN = 256; + static constexpr int kNumSubTiles = 1; + static constexpr int kLogitStride = kTileN; + static constexpr int kClusterM = 2; + static constexpr int kWarpSize = + fused_scaled_linear_cross_entropy::kWarpSize; + static constexpr int kConsumerThreads = 256; + static constexpr int kEpilogueThreads = 96; + static constexpr int kProducerRegisters = 88; + static constexpr int kConsumerRegisters = 208; + static constexpr int kMmaBarrierId = 9; + static constexpr int kEpilogueBarrierId = 10; + + using ClusterShape = Shape, _1, _1>; + using TiledMmaTraits = sm90::TiledMmaMSplitSm90< + Element, + Element, + ElementAccum, + kTileN, + 2, + GMMA::Major::K, + GMMA::Major::K, + Compute>; + using TiledMma = typename TiledMmaTraits::Type; + using AccumShape = Shape, Int>; + + using SmemTraitsA = + sm90::SmemLayoutKMajorSm90; + using SmemTraitsB = + sm90::SmemLayoutKMajorSm90; + using SmemLayoutA = typename SmemTraitsA::Staged; + using SmemLayoutB = typename SmemTraitsB::Staged; + using SmemStoreSingle = Layout< + Shape, Int>, + Stride, _1>>; + using SmemStore = Layout< + Shape, Int, Int>, + Stride, _1, Int>>; + + using PipelineTraits = sm90::MainloopPipelineSm90; + using MainloopPipeline = typename PipelineTraits::Type; + using PipelineState = typename PipelineTraits::State; + + static constexpr int kSmemAElements = SmemTraitsA::kStagedElements; + static constexpr int kSmemBElements = SmemTraitsB::kStagedElements; + static constexpr int kLogitElements = kTileM * kLogitStride; + static constexpr int kTmaTransBytes = + SmemTraitsA::kStageBytes + SmemTraitsB::kStageBytes; + + static_assert( + kProducerRegisters * 128 + kConsumerRegisters * 256 == 64512); +}; + + + +// ─────────────────────────────────────────────────────────────────────────── +// Shared memory +// ─────────────────────────────────────────────────────────────────────────── + +template +struct BackwardDxDwSmemSm90 { + using Traits = BackwardGemmTraitsSm90; + using Element = typename Traits::Element; + + alignas(16) typename Traits::MainloopPipeline::SharedStorage mainloop_pipeline; + alignas(8) std::uint64_t + dx_ready_epoch[kDxCommWarpsPerChannel][kDxRingStages]; + alignas(8) std::uint64_t + dx_consumed_barrier[kDxCommWarpsPerChannel][kDxRingStages]; + + alignas(1024) Element smem_a[Traits::kSmemAElements]; + alignas(1024) Element smem_b[Traits::kSmemBElements]; + alignas(1024) Element smem_store[Traits::kStoreElements]; + + CUTE_DEVICE Element* a_data() { return &smem_a[0]; } + CUTE_DEVICE Element* b_data() { return &smem_b[0]; } + CUTE_DEVICE Element* store_data() { return &smem_store[0]; } +}; + + +template +struct BackwardDxCluster2SmemSm90 { + using Traits = BackwardDxCluster2TraitsSm90; + using DwTraits = BackwardGemmTraitsSm90; + using Element = typename Traits::Element; + static constexpr int kDxStoreBytes = + Traits::SmemStore::kStagedElements * sizeof(float); + static constexpr int kDwStoreBytes = + DwTraits::kDwStoreElements * sizeof(Element); + static constexpr int kStoreBytes = + kDxStoreBytes > kDwStoreBytes ? kDxStoreBytes : kDwStoreBytes; + static constexpr int kStoreElements = kStoreBytes / sizeof(Element); + + alignas(16) typename Traits::MainloopPipeline::SharedStorage mainloop_pipeline; + alignas(8) std::uint64_t + dx_ready_epoch[kDxCommWarpsPerChannel][kDxRingStages]; + alignas(8) std::uint64_t + dx_consumed_barrier[kDxCommWarpsPerChannel][kDxRingStages]; + alignas(1024) Element smem_a[Traits::kSmemAElements]; + alignas(1024) Element smem_b0[Traits::kSmemBElements]; + alignas(1024) Element smem_store[kStoreElements]; + + CUTE_DEVICE Element* a_data() { return &smem_a[0]; } + CUTE_DEVICE Element* b_data() { return &smem_b0[0]; } + CUTE_DEVICE Element* b0_data() { return &smem_b0[0]; } + CUTE_DEVICE Element* store_data() { return &smem_store[0]; } + CUTE_DEVICE float* dx_store_data() { + return reinterpret_cast(&smem_store[0]); + } +}; + +template +struct BackwardDzHandoffSmemSm90 { + using Traits = BackwardDzHandoffTraitsSm90; + using Element = typename Traits::Element; + using ElementLogit = typename Traits::ElementLogit; + + alignas(16) typename Traits::MainloopPipeline::SharedStorage mainloop_pipeline; + alignas(8) std::uint64_t logits_ready; + alignas(8) std::uint64_t logits_free; + alignas(16) float row_scale[Traits::kTileM]; + alignas(16) float row_lse[Traits::kTileM]; + alignas(16) float row_entropy[ReturnEntropy ? Traits::kTileM : 1]; + alignas(16) float row_entropy_scale[ReturnEntropy ? Traits::kTileM : 1]; + alignas(16) int row_target[Traits::kTileM]; + + alignas(1024) Element smem_a[Traits::kSmemAElements]; + alignas(1024) Element smem_b[Traits::kSmemBElements]; + alignas(1024) ElementLogit smem_logits[Traits::kLogitElements]; + + CUTE_DEVICE Element* a_data() { return &smem_a[0]; } + CUTE_DEVICE Element* b_data() { return &smem_b[0]; } + CUTE_DEVICE ElementLogit* logit_row(int row) { + return &smem_logits[row * Traits::kLogitStride]; + } + CUTE_DEVICE Element* store_data() { + return reinterpret_cast(&smem_logits[0]); + } +}; + + + +// ─────────────────────────────────────────────────────────────────────────── +// Named barriers +// ─────────────────────────────────────────────────────────────────────────── + +template +CUTE_DEVICE void backward_consumer_barrier_sm90() { + using Config = BackwardGemmConfigSm90; + sm90::named_barrier_sync_sm90< + BarrierId, Config::kConsumerThreads, Compute>(); +} + +// Hands the aliased A/B and store-staging arenas from the dX phase to the dW +// phase across warp 0 and warps 4..11 only. __syncthreads() would drag the +// communication warps in and serialise the dX all-reduce tail against the dW +// GEMM, which is the whole point of keeping them separate. +template +CUTE_DEVICE void backward_compute_barrier_sm90() { + using Config = BackwardGemmConfigSm90; + sm90::named_barrier_sync_sm90< + Config::kComputeBarrierId, Config::kComputeBarrierThreads, Compute>(); +} + +CUTE_DEVICE void dx_ready_store_release( + std::uint64_t* ready, std::uint64_t epoch) { + cuda::atomic_ref flag(*ready); + flag.store(epoch, cuda::memory_order_release); +} + +CUTE_DEVICE void dx_ready_wait_acquire( + std::uint64_t* ready, std::uint64_t epoch) { + cuda::atomic_ref flag(*ready); + while (flag.load(cuda::memory_order_acquire) != epoch) { + __nanosleep(64); + } +} + +// ─────────────────────────────────────────────────────────────────────────── +// dZ phase +// ─────────────────────────────────────────────────────────────────────────── + +struct BackwardTileCoord { + int m_tile; + int n_tile; +}; + +// Flat n-fastest raster; CTA `cta` takes every `num_ctas`-th tile. +__host__ __device__ inline BackwardTileCoord backward_raster_tile( + int index, int num_n_tiles) { + BackwardTileCoord coord; + coord.m_tile = index / num_n_tiles; + coord.n_tile = index - coord.m_tile * num_n_tiles; + return coord; +} + +// Cluster peers own adjacent token tiles and therefore reuse one multicast W +// panel. Token-pair tiles advance fastest to match the winning dX schedule. +__host__ __device__ inline BackwardTileCoord backward_cluster_m_m_fast_tile( + int index, int num_m_pairs, int cluster_rank) { + BackwardTileCoord coord; + coord.m_tile = 2 * (index % num_m_pairs) + cluster_rank; + coord.n_tile = index / num_m_pairs; + return coord; +} + +// Cluster peers keep the same vocabulary tile while rank selects adjacent +// hidden tiles. Hidden-N pairs are the fastest-moving work dimension. +__host__ __device__ inline BackwardTileCoord backward_cluster_n_tile( + int index, int num_n_cluster_tiles, int cluster_rank) { + BackwardTileCoord coord = + backward_raster_tile(index, num_n_cluster_tiles); + coord.n_tile = 2 * coord.n_tile + cluster_rank; + return coord; +} + +__host__ __device__ inline BackwardTileCoord backward_cluster_n_m_fast_tile( + int index, int num_m_tiles, int cluster_rank) { + BackwardTileCoord coord; + coord.m_tile = index % num_m_tiles; + coord.n_tile = 2 * (index / num_m_tiles) + cluster_rank; + return coord; +} + +// ─────────────────────────────────────────────────────────────────────────── +// dX phase +// ─────────────────────────────────────────────────────────────────────────── + +template +struct DxPhaseSm90 { + using Traits = BackwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Element = typename Traits::Element; + static constexpr int kCompute = Compute; + static constexpr int kBarrierId = Config::kDxEpilogueBarrierId; + static_assert( + CommConfig::kTileElements == Traits::kTileM * Traits::kTileN, + "dX staging tile stride must match the WGMMA accumulator tile"); + + template < + bool ReturnEntropy, + class Pipeline, + class Smem, + class TmaDz, + class TmaWt> + CUTE_DEVICE static void produce( + Pipeline& pipe, + typename Traits::PipelineState& state, + Smem& smem, + TmaDz const& tma_dz, + TmaWt const& tma_wt, + BackwardGemmParamsSm90 const& params, + int padded_vocab, + int num_n_tiles, + int num_k_tiles, + int groups_per_wave, + int first_unit, + int stride_unit) { + if (cute::elect_one_sync() == 0) return; + + auto sA = make_tensor(make_smem_ptr(smem.a_data()), + typename Traits::SmemAKMajor::Staged{}); + auto sB = make_tensor(make_smem_ptr(smem.b_data()), + typename Traits::SmemBMnMajor::Staged{}); + + auto mDz = tma_dz.get_tma_tensor(make_shape( + static_cast(Config::kWaveRows), + static_cast(padded_vocab))); + // W as [hidden, vocab]: the dX B operand is MN-major. + auto mWt = tma_wt.get_tma_tensor(make_shape( + static_cast(params.hidden), + static_cast(params.local_vocab))); + + auto cta_dz = tma_dz.get_slice(Int<0>{}); + auto cta_wt = tma_wt.get_slice(Int<0>{}); + auto tAsA = cta_dz.partition_D(sA); + auto tBsB = cta_wt.partition_D(sB); + + for (int unit = first_unit; unit < groups_per_wave; + unit += stride_unit) { + DxTileGroup group = + dx_unit_to_group(unit, num_n_tiles, 0); + auto gDz = local_tile( + mDz, + make_tile(Int{}, Int{}), + make_coord(group.m_tile, _)); + auto tAgA = cta_dz.partition_S(gDz); + + for (int tile = 0; tile < group.num_tiles; ++tile) { + auto gWt = local_tile( + mWt, + make_tile(Int{}, Int{}), + make_coord(group.first_n_tile + tile, _)); + auto tBgB = cta_wt.partition_S(gWt); + + sm90::MainloopKLoopSm90::produce( + pipe, state, num_k_tiles, + [&](auto* barrier, int stage, int k_tile) { + copy(tma_dz.with(*barrier, 0), + tAgA(_, _, _, k_tile), tAsA(_, _, _, stage)); + copy(tma_wt.with(*barrier, 0), + tBgB(_, _, _, k_tile), tBsB(_, _, _, stage)); + }); + } + } + } + + // Consumer: GEMM, then store the complete FP32 group into this CTA's + // compact symmetric staging and publish one CTA-local ready arrival. + template + CUTE_DEVICE static int consume( + Pipeline& pipe, + typename Traits::PipelineState& read_state, + typename Traits::PipelineState& release_state, + Smem& smem, + DxReduceWorkspace const& comm, + int num_n_tiles, + int num_k_tiles, + int wave, + int groups_per_wave, + int first_unit, + int stride_unit, + int next_index) { + typename Traits::DxMmaTraits::Type tiled_mma; + int tid_in_mma = + static_cast(threadIdx.x) - Config::kFirstConsumerThread; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sA = make_tensor(make_smem_ptr(smem.a_data()), + typename Traits::SmemAKMajor::Staged{}); + auto sB = make_tensor(make_smem_ptr(smem.b_data()), + typename Traits::SmemBMnMajor::Staged{}); + auto tCsA = thr_mma.partition_A(sA); + auto tCsB = thr_mma.partition_B(sB); + + auto accum = partition_fragment_C( + tiled_mma, typename Traits::AccumShape{}); + auto coord_c = thr_mma.partition_C( + make_identity_tensor(typename Traits::AccumShape{})); + + for (int unit = first_unit; unit < groups_per_wave; + unit += stride_unit) { + int group_id = wave * groups_per_wave + unit; + DxTileGroup group = + dx_unit_to_group(unit, num_n_tiles, group_id); + DxCtaGroupSlot slot = + dx_cta_group_slot(next_index); + + // Only this CTA's comm warps can release this slot. The elected + // consumer waits for the prior pass, then releases all consumers. + if (slot.pass > 0) { + if (tid_in_mma == 0) { + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + cute::wait_barrier( + smem.dx_consumed_barrier[ + comm_warp][slot.stage], + (slot.pass - 1) & 1); + } + } + backward_consumer_barrier_sm90(); + } + + for (int tile = 0; tile < group.num_tiles; ++tile) { + CUTE_UNROLL + for (int i = 0; + i < decltype(cute::size(accum))::value; + ++i) { + asm volatile( + "mov.f32 %0, 0f00000000;" : "=f"(accum(i))); + } + sm90::MainloopKLoopSm90::consume( + pipe, read_state, release_state, num_k_tiles, + [&](int stage) { + warpgroup_fence_operand(accum); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, stage), + tCsB(_, _, _, stage), accum); + warpgroup_commit_batch(); + }, + [&]() { + warpgroup_fence_operand(accum); + }); + + // Materialize every FP32 accumulator element in the staging + // arena before communication consumes the tile. + std::size_t tile_base = dx_slot_offset( + static_cast(blockIdx.x), + slot.comm_warp, + slot.stage) + + static_cast(tile) * + CommConfig::kTileElements; + CUTE_UNROLL + for (int i = 0; i < decltype(cute::size(accum))::value; ++i) { + int row = get<0>(coord_c(i)); + int column = get<1>(coord_c(i)); + comm.partial[tile_base + + static_cast(row) * Traits::kTileN + + column] = accum(i); + } + } + + // The named barrier orders every scalar writer before the elected + // writer's CTA-release epoch. The comm warp publishes system-wide + // only after its matching acquire. + backward_consumer_barrier_sm90(); + if (tid_in_mma == 0) { + std::uint64_t ready_epoch = + static_cast(group_id) + 1; + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + dx_ready_store_release( + &smem.dx_ready_epoch[comm_warp][slot.stage], + ready_epoch); + } + } + ++next_index; + } + return next_index; + } + + struct alignas(8) Bfloat16x4 { + __nv_bfloat162 low; + __nv_bfloat162 high; + }; + +}; + +// ─────────────────────────────────────────────────────────────────────────── +// dW phase — rank local, unchanged from the single-GPU reference. +// ─────────────────────────────────────────────────────────────────────────── + +template < + int Compute = 90, + class Traits_ = BackwardGemmTraitsSm90> +struct DwPhaseSm90 { + using Traits = Traits_; + using Config = typename Traits::Config; + using Launch = BackwardGemmLaunchSm90; + static constexpr int kCompute = Compute; + static constexpr int kStoreTileN = Traits::kDwStoreTileN; + static constexpr int kStoreStages = Traits::kDwStoreStages; + static constexpr int kNumSubTiles = Traits::kTileN / kStoreTileN; + static constexpr int kBarrierId = Config::kDwStoreBarrierId; + + template < + bool ReturnEntropy, + class Pipeline, + class Smem, + class TmaDzt, + class TmaXt> + CUTE_DEVICE static void produce( + Pipeline& pipe, + typename Traits::PipelineState& state, + Smem& smem, + TmaDzt const& tma_dzt, + TmaXt const& tma_xt, + BackwardGemmParamsSm90 const& params, + int padded_vocab, + int wave, + int num_n_tiles, + int num_k_tiles, + int num_tiles, + int first_index, + int stride_index, + int cluster_n_rank = -1, + std::uint16_t multicast_mask = 0) { + if (cute::elect_one_sync() == 0) return; + + auto sA = make_tensor(make_smem_ptr(smem.a_data()), + typename Traits::SmemAMnMajor::Staged{}); + auto sB = make_tensor(make_smem_ptr(smem.b_data()), + typename Traits::SmemBMnMajor::Staged{}); + + // dZ^T as [vocab, wave_rows] and X^T as [hidden, tokens]. + auto mDzt = tma_dzt.get_tma_tensor(make_shape( + static_cast(padded_vocab), + static_cast(Config::kWaveRows))); + auto mXt = tma_xt.get_tma_tensor(make_shape( + static_cast(params.hidden), + static_cast(params.tokens))); + + auto cta_dzt = + tma_dzt.get_slice(cluster_n_rank < 0 ? 0 : cluster_n_rank); + auto cta_xt = tma_xt.get_slice(Int<0>{}); + auto tAsA = cta_dzt.partition_D(sA); + auto tBsB = cta_xt.partition_D(sB); + + // The wave's token rows are the K extent of this GEMM. + int k_base = wave * (Config::kWaveRows / Traits::kTileK); + for (int t = 0; t < num_tiles; ++t) { + int work = first_index + t * stride_index; + BackwardTileCoord tile = cluster_n_rank < 0 + ? backward_raster_tile(work, num_n_tiles) + : backward_cluster_n_tile( + work, num_n_tiles, cluster_n_rank); + auto gDzt = local_tile( + mDzt, + make_tile(Int{}, Int{}), + make_coord(tile.m_tile, _)); + auto gXt = local_tile( + mXt, + make_tile(Int{}, Int{}), + make_coord(tile.n_tile, _)); + auto tAgA = cta_dzt.partition_S(gDzt); + auto tBgB = cta_xt.partition_S(gXt); + + sm90::MainloopKLoopSm90::produce( + pipe, state, num_k_tiles, + [&](auto* barrier, int stage, int k_tile) { + copy(tma_dzt.with(*barrier, multicast_mask), + tAgA(_, _, _, k_tile), tAsA(_, _, _, stage)); + copy(tma_xt.with(*barrier, 0), + tBgB(_, _, _, k_base + k_tile), tBsB(_, _, _, stage)); + }); + } + } + + template < + bool ReturnEntropy, + class Pipeline, + class Accumulator, + class Smem, + class TmaDwStore, + class TmaDwAdd> + CUTE_DEVICE static void consume_with_accumulator( + Pipeline& pipe, + typename Traits::PipelineState& read_state, + typename Traits::PipelineState& release_state, + Accumulator& accum, + Smem& smem, + TmaDwStore const& tma_store, + TmaDwAdd const& tma_add, + BackwardGemmParamsSm90 const& params, + int wave, + int num_n_tiles, + int num_k_tiles, + int num_tiles, + int first_index, + int stride_index, + int cluster_n_rank = -1) { + typename Traits::DwMmaTraits::Type tiled_mma; + int tid_in_mma = + static_cast(threadIdx.x) - Config::kFirstConsumerThread; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sA = make_tensor(make_smem_ptr(smem.a_data()), + typename Traits::SmemAMnMajor::Staged{}); + auto sB = make_tensor(make_smem_ptr(smem.b_data()), + typename Traits::SmemBMnMajor::Staged{}); + auto tCsA = thr_mma.partition_A(sA); + auto tCsB = thr_mma.partition_B(sB); + + using ExpectedAccumulator = decltype(partition_fragment_C( + tiled_mma, typename Traits::AccumShape{})); + static_assert(cute::is_same_v); + + int valid_n_tiles = ceil_div(params.hidden, Traits::kTileN); + for (int t = 0; t < num_tiles; ++t) { + CUTE_UNROLL + for (int i = 0; i < decltype(cute::size(accum))::value; ++i) { + asm volatile( + "mov.f32 %0, 0f00000000;" : "=f"(accum(i))); + } + + sm90::MainloopKLoopSm90::consume( + pipe, read_state, release_state, num_k_tiles, + [&](int stage) { + warpgroup_fence_operand(accum); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, stage), + tCsB(_, _, _, stage), accum); + warpgroup_commit_batch(); + }, + [&]() { warpgroup_fence_operand(accum); }); + + + int work = first_index + t * stride_index; + BackwardTileCoord tile = cluster_n_rank < 0 + ? backward_raster_tile(work, num_n_tiles) + : backward_cluster_n_tile( + work, num_n_tiles, cluster_n_rank); + bool valid_tile = tile.n_tile < valid_n_tiles; + + bool issuer = (tid_in_mma / Traits::kWarpSize) == 0; + auto sStore = make_tensor(make_smem_ptr(smem.store_data()), + typename Traits::template DwSmemStore< + Traits::kDwStoreStages>::Staged{}); + auto shape_dw = make_shape( + static_cast(params.local_vocab), + static_cast(params.hidden)); + auto mStore = tma_store.get_tma_tensor(shape_dw); + auto cta_store = tma_store.get_slice(Int<0>{}); + auto tSsS = cta_store.partition_S(sStore); + auto tSgS = cta_store.partition_D(local_tile( + mStore, + make_tile(Int{}, Int{}), + make_coord(_, _))); + auto mAdd = tma_add.get_tma_tensor(shape_dw); + auto cta_add = tma_add.get_slice(Int<0>{}); + auto tAsA = cta_add.partition_S(sStore); + auto tAgA = cta_add.partition_D(local_tile( + mAdd, + make_tile(Int{}, Int{}), + make_coord(_, _))); + + typename Traits::DwStoreMmaTraits::Type store_mma; + auto store_thr = store_mma.get_slice(tid_in_mma); + auto tDsDFrag = store_thr.partition_C(sStore); + auto d_frag = make_tensor( + shape(tDsDFrag(_, _, _, Int<0>{}))); + constexpr int kFragElements = + decltype(cute::size(d_frag))::value; + static_assert( + kFragElements * kNumSubTiles == + decltype(cute::size(accum))::value, + "dW store fragments must exactly cover the accumulator"); + + CUTE_UNROLL + for (int sub = 0; sub < kNumSubTiles; ++sub) { + int stage = sub % kStoreStages; + if (sub >= kStoreStages) { + if (issuer) { + wait_tma_store< + kStoreStages - 1>(); + } + backward_consumer_barrier_sm90(); + } + CUTE_UNROLL + for (int j = 0; j < kFragElements; ++j) { + d_frag(j) = static_cast( + accum(sub * kFragElements + j)); + } + copy(d_frag, tDsDFrag(_, _, _, stage)); + tma_store_fence(); + backward_consumer_barrier_sm90(); + if (valid_tile && issuer && cute::elect_one_sync()) { + int n = tile.n_tile * kNumSubTiles + sub; + // Wave 0 initializes dW; later waves accumulate with + // SM90_TMA_REDUCE_ADD in the same tile order. + if (wave == 0) { + copy(tma_store, + tSsS(_, _, _, stage), + tSgS(_, _, _, tile.m_tile, n)); + } else { + copy(tma_add, + tAsA(_, _, _, stage), + tAgA(_, _, _, tile.m_tile, n)); + } + tma_store_arrive(); + } + } + + if (issuer) { + wait_tma_store_complete(); + asm volatile("fence.proxy.async.global;" ::: "memory"); + } + backward_consumer_barrier_sm90(); + } + } + + template < + bool ReturnEntropy, + class Pipeline, + class Smem, + class TmaDwStore, + class TmaDwAdd> + CUTE_DEVICE static void consume( + Pipeline& pipe, + typename Traits::PipelineState& read_state, + typename Traits::PipelineState& release_state, + Smem& smem, + TmaDwStore const& tma_store, + TmaDwAdd const& tma_add, + BackwardGemmParamsSm90 const& params, + int wave, + int num_n_tiles, + int num_k_tiles, + int num_tiles, + int first_index, + int stride_index, + int cluster_n_rank = -1) { + typename Traits::DwMmaTraits::Type tiled_mma; + auto accum = partition_fragment_C( + tiled_mma, typename Traits::AccumShape{}); + consume_with_accumulator( + pipe, + read_state, + release_state, + accum, + smem, + tma_store, + tma_add, + params, + wave, + num_n_tiles, + num_k_tiles, + num_tiles, + first_index, + stride_index, + cluster_n_rank); + } +}; + +template < + class TmaDzLoad, + class TmaWt, + class TmaDzt, + class TmaXt, + class TmaDwStore, + class TmaDwAdd> +struct BackwardDxDwTmaBundleSm90 { + TmaDzLoad dz_load; + TmaWt wt; + TmaDzt dzt; + TmaXt xt; + TmaDwStore dw_store; + TmaDwAdd dw_add; +}; + +template < + bool ReturnEntropy, + int Compute, + class TmaX, + class TmaW, + class TmaDzStore> +__global__ __launch_bounds__(BackwardGemmConfigSm90::kNumThreads, 1) +void backward_dz_handoff_wave_kernel_sm90( + __grid_constant__ const TmaX tma_x, + __grid_constant__ const TmaW tma_w, + __grid_constant__ const TmaDzStore tma_dz_store, + __grid_constant__ const BackwardGemmParamsSm90 params, + int wave) { + using Traits = BackwardDzHandoffTraitsSm90; + using Config = typename Traits::Config; + using Launch = BackwardGemmLaunchSm90; + using Smem = BackwardDzHandoffSmemSm90; + using PipelineTraits = typename Traits::PipelineTraits; + using Registers = sm90::WarpSpecializedRegistersSm90< + Traits::kProducerRegisters, + Traits::kConsumerRegisters, + Compute>; + + extern __shared__ char raw_smem[]; + Smem& smem = *reinterpret_cast(raw_smem); + int warp_group = cutlass::canonical_warp_group_idx(); + int warp = cutlass::canonical_warp_idx_sync(); + int cluster_rank = + static_cast(cute::block_rank_in_cluster()); + int cluster_idx = static_cast(blockIdx.z); + int num_clusters = static_cast(gridDim.z); + int num_n_cluster_tiles = ceil_div( + Launch::padded_vocab(params.local_vocab), + Traits::kClusterM * Traits::kTileN); + int total_tiles = + Config::kMTilesPerWave * num_n_cluster_tiles; + int tiles = total_tiles > cluster_idx + ? ceil_div(total_tiles - cluster_idx, num_clusters) + : 0; + + PipelineTraits::init_barriers( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + typename Traits::ClusterShape{}); + if (threadIdx.x == 0) { + cute::initialize_barrier(smem.logits_ready, 1); + cute::initialize_barrier(smem.logits_free, 1); + } + sm90::cluster_pipeline_init_sm90(); + + if (warp_group == 0) { + Registers::producer(); + if (warp == 3) { + auto state = PipelineTraits::producer_start_state(); + auto pipe = PipelineTraits::make_producer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + threadIdx.x == 3 * Traits::kWarpSize, + typename Traits::ClusterShape{}); + cute::prefetch_tma_descriptor(tma_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_w.get_tma_descriptor()); + if (cute::elect_one_sync()) { + auto sA = make_tensor( + make_smem_ptr(smem.a_data()), + typename Traits::SmemLayoutA{}); + auto sB = make_tensor( + make_smem_ptr(smem.b_data()), + typename Traits::SmemLayoutB{}); + auto mX = tma_x.get_tma_tensor(make_shape( + static_cast(params.tokens), + static_cast(params.hidden))); + auto mW = tma_w.get_tma_tensor(make_shape( + static_cast(params.local_vocab), + static_cast(params.hidden))); + auto cta_x = tma_x.get_slice(Int<0>{}); + auto cta_w = tma_w.get_slice(Int<0>{}); + auto tAsA = cta_x.partition_D(sA); + auto tBsB = cta_w.partition_D(sB); + int m_base = wave * Config::kMTilesPerWave; + + for (int t = 0; t < tiles; ++t) { + BackwardTileCoord tile = backward_cluster_n_m_fast_tile( + cluster_idx + t * num_clusters, + Config::kMTilesPerWave, + cluster_rank); + auto gX = local_tile( + mX, + make_tile( + Int{}, + Int{}), + make_coord(m_base + tile.m_tile, _)); + auto gW = local_tile( + mW, + make_tile( + Int{}, + Int{}), + make_coord(tile.n_tile, _)); + auto tAgA = cta_x.partition_S(gX); + auto tBgB = cta_w.partition_S(gW); + sm90::MainloopKLoopSm90::produce( + pipe, + state, + Launch::num_dz_k_tiles(params.hidden), + [&](auto* barrier, int stage, int k_tile) { + copy( + tma_x.with(*barrier, 0), + tAgA(_, _, _, k_tile), + tAsA(_, _, _, stage)); + copy( + tma_w.with(*barrier, 0), + tBgB(_, _, _, k_tile), + tBsB(_, _, _, stage)); + }); + } + pipe.producer_tail(state); + } + } else if (warp < 3) { + int epi_tid = static_cast(threadIdx.x); + auto sStore = make_tensor( + make_smem_ptr(smem.store_data()), + typename Traits::SmemStore{}); + auto mDz = tma_dz_store.get_tma_tensor(make_shape( + static_cast(Config::kWaveRows), + static_cast( + Launch::padded_vocab(params.local_vocab)))); + auto cta_dz = tma_dz_store.get_slice(Int<0>{}); + auto tSsS = cta_dz.partition_S(sStore); + auto tSgS = cta_dz.partition_D(local_tile( + mDz, + make_tile( + Int{}, + Int{}), + make_coord(_, _))); + float exp_scale = + params.inverse_temperature * kBackwardLog2E; + + for (int t = 0; t < tiles; ++t) { + if (epi_tid == 0) { + cute::wait_barrier(smem.logits_ready, t & 1); + } + sm90::named_barrier_sync_sm90< + Traits::kEpilogueBarrierId, + Traits::kEpilogueThreads, + Compute>(); + BackwardTileCoord tile = backward_cluster_n_m_fast_tile( + cluster_idx + t * num_clusters, + Config::kMTilesPerWave, + cluster_rank); + + for (int row = epi_tid; + row < Traits::kTileM; + row += Traits::kEpilogueThreads) { + int token = + (wave * Config::kMTilesPerWave + tile.m_tile) * + Traits::kTileM + + row; + float scale = 0.0f; + float exp_bias = 0.0f; + float correction_bias = 0.0f; + float correction_slope = 0.0f; + int target_local = -1; + if (token < params.tokens) { + std::int64_t target_id = params.target[token]; + if (target_id != params.ignore_index) { + float lse = params.lse[token]; + scale = params.grad_output[token]; + exp_bias = -lse * kBackwardLog2E; + if constexpr (ReturnEntropy) { + float entropy = params.entropy[token]; + float entropy_scale = + params.entropy_grad[token]; + correction_bias = fmaf( + lse - entropy, + entropy_scale, + scale); + correction_slope = + -params.inverse_temperature * + entropy_scale; + } + std::int64_t local = + target_id - params.vocab_start; + if (local >= 0 && + local < static_cast( + params.local_vocab)) { + target_local = static_cast(local); + } + } + } + smem.row_scale[row] = scale; + smem.row_lse[row] = exp_bias; + smem.row_target[row] = target_local; + if constexpr (ReturnEntropy) { + smem.row_entropy[row] = correction_bias; + smem.row_entropy_scale[row] = correction_slope; + } + } + sm90::named_barrier_sync_sm90< + Traits::kEpilogueBarrierId, + Traits::kEpilogueThreads, + Compute>(); + + int vocab_base = tile.n_tile * Traits::kTileN; + int lane = epi_tid & (Traits::kWarpSize - 1); + CUTE_UNROLL + for (int sub = 0; + sub < Traits::kNumSubTiles; + ++sub) { + for (int row = warp; + row < Traits::kTileM; + row += 3) { + float scale = smem.row_scale[row]; + float exp_bias = smem.row_lse[row]; + int target = smem.row_target[row]; + float correction_bias = 0.0f; + float correction_slope = 0.0f; + if constexpr (ReturnEntropy) { + correction_bias = smem.row_entropy[row]; + correction_slope = + smem.row_entropy_scale[row]; + } + CUTE_UNROLL + for (int j = lane; + j < Traits::kStoreTileN; + j += Traits::kWarpSize) { + int column = + vocab_base + + sub * Traits::kStoreTileN + j; + float value = 0.0f; + if (column < params.local_vocab) { + float logit = static_cast( + smem.logit_row(row)[ + sub * Traits::kStoreTileN + j]); + if constexpr (ReturnEntropy) { + if (scale != 0.0f || + correction_slope != 0.0f) { + float probability = + backward_exp2_sm90(fmaf( + logit, + exp_scale, + exp_bias)); + value = probability * fmaf( + logit, + correction_slope, + correction_bias); + } + } else if (scale != 0.0f) { + value = backward_exp2_sm90(fmaf( + logit, + exp_scale, + exp_bias)) * + scale; + } + if (column == target) value -= scale; + value *= params.inverse_temperature; + } + sStore(row, j, sub) = + static_cast(value); + } + } + tma_store_fence(); + sm90::named_barrier_sync_sm90< + Traits::kEpilogueBarrierId, + Traits::kEpilogueThreads, + Compute>(); + if (warp == 0 && cute::elect_one_sync()) { + copy( + tma_dz_store, + tSsS(_, _, _, sub), + tSgS( + _, _, _, tile.m_tile, + tile.n_tile * + Traits::kNumSubTiles + + sub)); + tma_store_arrive(); + } + } + sm90::named_barrier_sync_sm90< + Traits::kEpilogueBarrierId, + Traits::kEpilogueThreads, + Compute>(); + if (warp == 0) { + wait_tma_store_complete(); + } + sm90::named_barrier_sync_sm90< + Traits::kEpilogueBarrierId, + Traits::kEpilogueThreads, + Compute>(); + if (epi_tid == 0) { + cute::arrive_barrier(smem.logits_free); + } + } + } + } else { + Registers::consumer(); + typename Traits::TiledMma tiled_mma; + int tid_in_mma = + static_cast(threadIdx.x) - Config::kFirstConsumerThread; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + auto sA = make_tensor( + make_smem_ptr(smem.a_data()), + typename Traits::SmemLayoutA{}); + auto sB = make_tensor( + make_smem_ptr(smem.b_data()), + typename Traits::SmemLayoutB{}); + auto tCsA = thr_mma.partition_A(sA); + auto tCsB = thr_mma.partition_B(sB); + auto accum = partition_fragment_C( + tiled_mma, typename Traits::AccumShape{}); + auto coord_c = thr_mma.partition_C( + make_identity_tensor(typename Traits::AccumShape{})); + typename Traits::PipelineState read_state; + typename Traits::PipelineState release_state; + auto pipe = PipelineTraits::make_consumer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + typename Traits::ClusterShape{}); + + for (int t = 0; t < tiles; ++t) { + CUTE_UNROLL + for (int i = 0; i < decltype(cute::size(accum))::value; ++i) { + asm volatile( + "mov.f32 %0, 0f00000000;" : "=f"(accum(i))); + } + sm90::MainloopKLoopSm90::consume( + pipe, + read_state, + release_state, + Launch::num_dz_k_tiles(params.hidden), + [&](int stage) { + warpgroup_fence_operand(accum); + warpgroup_arrive(); + gemm( + tiled_mma, + tCsA(_, _, _, stage), + tCsB(_, _, _, stage), + accum); + warpgroup_commit_batch(); + }, + [&]() { warpgroup_fence_operand(accum); }); + + if (t > 0) { + if (tid_in_mma == 0) { + cute::wait_barrier( + smem.logits_free, + (t - 1) & 1); + } + sm90::named_barrier_sync_sm90< + Traits::kMmaBarrierId, + Traits::kConsumerThreads, + Compute>(); + } + CUTE_UNROLL + for (int i = 0; i < decltype(cute::size(accum))::value; ++i) { + int row = get<0>(coord_c(i)); + int column = get<1>(coord_c(i)); + smem.logit_row(row)[column] = + static_cast(accum(i)); + } + cutlass::arch::fence_view_async_shared(); + asm volatile("fence.acq_rel.cta;" ::: "memory"); + sm90::named_barrier_sync_sm90< + Traits::kMmaBarrierId, + Traits::kConsumerThreads, + Compute>(); + if (tid_in_mma == 0) { + cute::arrive_barrier(smem.logits_ready); + } + } + } + + __syncthreads(); + sm90::cluster_exit_sm90(); +} + +template < + class TmaDz, + class TmaWt, + class TmaDxStore, + class TmaDzt, + class TmaXt, + class TmaDwStore, + class TmaDwAdd> +struct BackwardClusterCombinedTmaBundleSm90 { + TmaDz dz; + TmaWt wt; + TmaDxStore dx_store; + TmaDzt dzt; + TmaXt xt; + TmaDwStore dw_store; + TmaDwAdd dw_add; +}; + +template < + int DxSplitK, + int Compute, + class CommConfig, + class Bundle> +__global__ __launch_bounds__( + BackwardDxCluster2TraitsSm90::kNumThreads, 1) +void backward_dx_cluster2_gemm_wave_kernel_sm90( + __grid_constant__ const Bundle tma, + __grid_constant__ const DxReduceWorkspace comm, + __grid_constant__ const liger_cute::detail::NvlsReduceView mapping, + __grid_constant__ const BackwardGemmParamsSm90 params, + int wave) { + static_assert(DxSplitK == 1 || DxSplitK == 2); + using Traits = BackwardDxCluster2TraitsSm90; + using DwTraits = BackwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Launch = BackwardGemmLaunchSm90; + using Smem = BackwardDxCluster2SmemSm90; + using PipelineTraits = typename Traits::PipelineTraits; + using ClusterShape = typename Traits::ClusterShape; + using Registers = sm90::WarpSpecializedRegistersSm90< + DxSplitK == 2 ? 72 : 56, + DxSplitK == 2 ? 216 : 224, + Compute>; + static_assert(Traits::kStages == 4); + static_assert(DwTraits::kStages == 4); + static_assert(DwTraits::kDwStoreTileN == 64); + static_assert(DwTraits::kDwStoreStages == 2); + static_assert(Traits::kTileM == DwTraits::kTileM); + static_assert(Traits::kTileN == DwTraits::kTileN); + static_assert(Traits::kTileK == DwTraits::kTileK); + static_assert(Traits::kTmaTransBytes == DwTraits::kTmaTransBytes); + static_assert(cute::is_same_v< + typename Traits::PipelineState, + typename DwTraits::PipelineState>); + + extern __shared__ char raw_smem[]; + Smem& smem = *reinterpret_cast(raw_smem); + int warp_group = cutlass::canonical_warp_group_idx(); + constexpr int kDwKTiles = Launch::num_dw_k_tiles(); + + PipelineTraits::init_barriers( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + if (threadIdx.x == 0) { + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + CUTE_UNROLL + for (int stage = 0; stage < kDxRingStages; ++stage) { + cute::initialize_barrier( + smem.dx_ready_epoch[comm_warp][stage], 1); + cute::initialize_barrier( + smem.dx_consumed_barrier[comm_warp][stage], 1); + } + } + } + sm90::cluster_pipeline_init_sm90(); + + if (warp_group == 0) { + Registers::producer(); + int warp = cutlass::canonical_warp_idx_sync(); + if (warp == 0) { + int cluster_rank = static_cast(cute::block_rank_in_cluster()); + int cluster_idx = static_cast(blockIdx.z); + int num_clusters = static_cast(gridDim.z); + int num_n_tiles = Launch::num_dx_n_tiles(params.hidden); + int num_k_tiles = + Launch::num_dx_k_tiles(params.local_vocab) / DxSplitK; + constexpr int num_m_pairs = + Config::kMTilesPerWave / Traits::kClusterM; + int total_cluster_tiles = + num_m_pairs * num_n_tiles * DxSplitK; + auto state = PipelineTraits::producer_start_state(); + auto pipe = PipelineTraits::make_producer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + threadIdx.x == 0, + ClusterShape{}); + cute::prefetch_tma_descriptor(tma.dz.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma.wt.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma.dx_store.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma.dzt.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma.xt.get_tma_descriptor()); + if (cute::elect_one_sync()) { + auto sA = make_tensor( + make_smem_ptr(smem.a_data()), + typename Traits::SmemA::Staged{}); + auto sB0 = make_tensor( + make_smem_ptr(smem.b0_data()), + typename Traits::SmemB::Staged{}); + auto mDz = tma.dz.get_tma_tensor(make_shape( + static_cast(Config::kWaveRows), + static_cast( + Launch::padded_vocab(params.local_vocab)))); + auto mWt = tma.wt.get_tma_tensor(make_shape( + static_cast(params.hidden), + static_cast(params.local_vocab))); + auto cta_dz = tma.dz.get_slice(Int<0>{}); + auto cta_wt = tma.wt.get_slice(cluster_rank); + auto tAsA = cta_dz.partition_D(sA); + auto tB0sB = cta_wt.partition_D(sB0); + std::uint16_t mcast_mask = + sm90::cluster_multicast_mask_m_sm90< + ClusterShape, Compute>(); + + for (int index = cluster_idx; + index < total_cluster_tiles; + index += num_clusters) { + int logical_index = index / DxSplitK; + int split_k = index % DxSplitK; + int first_k_tile = split_k * num_k_tiles; + BackwardTileCoord tile = + backward_cluster_m_m_fast_tile( + logical_index, num_m_pairs, cluster_rank); + auto gDz = local_tile( + mDz, + make_tile( + Int{}, + Int{}), + make_coord(tile.m_tile, _)); + auto gWt0 = local_tile( + mWt, + make_tile( + Int{}, + Int{}), + make_coord( + Traits::kNumFragments * tile.n_tile, _)); + auto tAgA = cta_dz.partition_S(gDz); + auto tB0gB = cta_wt.partition_S(gWt0); + // dZ is loaded per CTA; only the shared W^T panel is + // multicast to the adjacent-token cluster peer. + sm90::MainloopKLoopSm90::produce( + pipe, state, num_k_tiles, + [&](auto* barrier, int stage, int k_tile) { + copy(tma.dz.with(*barrier, 0), + tAgA(_, _, _, first_k_tile + k_tile), + tAsA(_, _, _, stage)); + copy(tma.wt.with(*barrier, mcast_mask), + tB0gB(_, _, _, first_k_tile + k_tile), + tB0sB(_, _, _, stage)); + }); + } + } + int dw_n_tiles = Launch::num_dw_n_tiles(params.hidden); + int dw_n_pairs = ceil_div(dw_n_tiles, Traits::kClusterM); + int dw_total_units = + Launch::num_dw_m_tiles(params.local_vocab) * dw_n_pairs; + int dw_units = dw_total_units > cluster_idx + ? ceil_div(dw_total_units - cluster_idx, num_clusters) + : 0; + std::uint16_t dzt_mcast_mask = + sm90::cluster_multicast_mask_m_sm90< + ClusterShape, Compute>(); + DwPhaseSm90::template produce( + pipe, + state, + smem, + tma.dzt, + tma.xt, + params, + Launch::padded_vocab(params.local_vocab), + wave, + dw_n_pairs, + kDwKTiles, + dw_units, + cluster_idx, + num_clusters, + cluster_rank, + dzt_mcast_mask); + if (cute::elect_one_sync()) pipe.producer_tail(state); + } else if ( + warp >= CommConfig::kFirstCommWarp && + warp <= CommConfig::kLastCommWarp) { + int cluster_rank = static_cast(cute::block_rank_in_cluster()); + int cluster_idx = static_cast(blockIdx.z); + int num_clusters = static_cast(gridDim.z); + int num_n_tiles = Launch::num_dx_n_tiles(params.hidden); + constexpr int num_m_pairs = + Config::kMTilesPerWave / Traits::kClusterM; + int total_cluster_tiles = + num_m_pairs * num_n_tiles * DxSplitK; + int cta = cluster_idx * Traits::kClusterM + cluster_rank; + int comm_warp = warp - CommConfig::kFirstCommWarp; + constexpr std::size_t kMessageElements = CommConfig::kTileElements; + constexpr std::size_t kSegmentElements = + kMessageElements / kDxCommWarpsPerChannel; + int local_index = 0; + for (int index = cluster_idx; + index < total_cluster_tiles; + index += num_clusters, ++local_index) { + int stage = local_index % CommConfig::kNumStages; + int pass = local_index / CommConfig::kNumStages; + if (static_cast(threadIdx.x) % Traits::kWarpSize == 0) { + cute::wait_barrier( + smem.dx_ready_epoch[comm_warp][stage], + pass & 1); + } + __syncwarp(); + + std::size_t base = + dx_slot_offset(cta, 0, stage); + std::size_t segment_begin = + static_cast(comm_warp) * + kSegmentElements; + BackwardTileCoord tile = backward_cluster_m_m_fast_tile( + index / DxSplitK, num_m_pairs, cluster_rank); + std::size_t durable_tile = + static_cast(tile.m_tile) * num_n_tiles + + tile.n_tile; + std::size_t tiles_per_wave = + static_cast(Config::kMTilesPerWave) * + num_n_tiles; + std::size_t packed_tile_elements = + kMessageElements / mapping.size; + std::size_t packed_segment_elements = + kSegmentElements / mapping.size; + float* durable = mapping.reduced_shard + + (static_cast(wave) * tiles_per_wave + + durable_tile) * packed_tile_elements + + static_cast(comm_warp) * + packed_segment_elements; + + if (mapping.size == 1) { + constexpr std::size_t kSegmentVectors = + kSegmentElements / 4; + for (std::size_t vector = + static_cast(threadIdx.x) % + Traits::kWarpSize; + vector < kSegmentVectors; + vector += Traits::kWarpSize) { + float4 value = + reinterpret_cast( + comm.partial + base + segment_begin)[vector]; + if constexpr (DxSplitK == 1) { + reinterpret_cast(durable)[vector] = + value; + } else { + float* output = durable + 4 * vector; + atomicAdd(output + 0, value.x); + atomicAdd(output + 1, value.y); + atomicAdd(output + 2, value.z); + atomicAdd(output + 3, value.w); + } + } + } else { + std::size_t ready_offset = + dx_sync_offset( + cta, + comm_warp, + stage, + kDxReadyPhase, + mapping.size); + std::size_t complete_offset = + dx_sync_offset( + cta, + comm_warp, + stage, + kDxCompletePhase, + mapping.size); + std::uint64_t epoch = + dx_epoch_base(comm) | + (static_cast(wave + 1) << 16) | + static_cast(pass + 1); + liger_cute::detail::nvls_barrier_warp( + comm.sync + ready_offset, + mapping.multicast_sync + ready_offset, + mapping.rank, + mapping.size, + epoch); + if constexpr (DxSplitK == 1) { + liger_cute::detail::nvls_sum_reduce_scatter_warp( + durable, + mapping.multicast_partial + + base + segment_begin, + kSegmentElements, + mapping.rank, + mapping.size); + } else { + liger_cute::detail::nvls_sum_reduce_scatter_warp( + mapping.multicast_reduced + + base + segment_begin, + mapping.multicast_partial + + base + segment_begin, + kSegmentElements, + mapping.rank, + mapping.size); + // Re-read the lane and accumulate one scalar at a time + // so no vector or lane value lives across the atomics. + for (std::size_t element = + static_cast(threadIdx.x) % + Traits::kWarpSize; + element < packed_segment_elements; + element += Traits::kWarpSize) { + atomicAdd( + durable + element, + mapping.multicast_reduced[ + base + segment_begin + element]); + } + __syncwarp(); + } + liger_cute::detail::nvls_barrier_warp( + comm.sync + complete_offset, + mapping.multicast_sync + complete_offset, + mapping.rank, + mapping.size, + epoch); + } + __syncwarp(); + if (static_cast(threadIdx.x) % Traits::kWarpSize == 0) { + cute::arrive_barrier( + smem.dx_consumed_barrier[ + comm_warp][stage]); + } + } + } + } else { + Registers::consumer(); + int cluster_rank = static_cast(cute::block_rank_in_cluster()); + int cluster_idx = static_cast(blockIdx.z); + int num_clusters = static_cast(gridDim.z); + int num_n_tiles = Launch::num_dx_n_tiles(params.hidden); + int num_k_tiles = + Launch::num_dx_k_tiles(params.local_vocab) / DxSplitK; + constexpr int num_m_pairs = + Config::kMTilesPerWave / Traits::kClusterM; + int total_cluster_tiles = + num_m_pairs * num_n_tiles * DxSplitK; + int cta = cluster_idx * Traits::kClusterM + cluster_rank; + typename Traits::PipelineState read_state; + typename Traits::PipelineState release_state; + auto pipe = PipelineTraits::make_consumer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + auto accum0 = partition_fragment_C( + typename Traits::TiledMma{}, typename Traits::AccumShape{}); + using DwMma = typename DwTraits::DwMmaTraits::Type; + using DwAccumulator = decltype(partition_fragment_C( + *static_cast(nullptr), + typename DwTraits::AccumShape{})); + static_assert(cute::is_same_v< + typename Traits::AccumShape, + typename DwTraits::AccumShape>); + static_assert(cute::is_same_v); + { + typename Traits::TiledMma tiled_mma; + int tid_in_mma = + static_cast(threadIdx.x) - Config::kFirstConsumerThread; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + auto sA = make_tensor( + make_smem_ptr(smem.a_data()), + typename Traits::SmemA::Staged{}); + auto sB0 = make_tensor( + make_smem_ptr(smem.b0_data()), + typename Traits::SmemB::Staged{}); + auto tCsA = thr_mma.partition_A(sA); + auto tCsB0 = thr_mma.partition_B(sB0); + auto sStore = make_tensor( + make_smem_ptr(smem.dx_store_data()), + typename Traits::SmemStore::Staged{}); + auto mStore = tma.dx_store.get_tma_tensor(make_shape( + static_cast( + kMaxDxResidentCtas * kDxCommWarpsPerChannel * + CommConfig::kNumStages * CommConfig::kTilesPerReduce * + Traits::kTileM), + static_cast(Traits::kTileN))); + auto cta_store = tma.dx_store.get_slice(Int<0>{}); + auto tSsS = cta_store.partition_S(sStore); + auto tSgS = cta_store.partition_D(local_tile( + mStore, + make_tile( + Int{}, Int{}), + make_coord(_, _))); + typename Traits::StoreMmaTraits::Type store_mma; + auto store_thread = store_mma.get_slice(tid_in_mma); + auto tDsD = store_thread.partition_C(sStore); + auto output_fragment = + make_tensor(shape(tDsD(_, _, _, Int<0>{}))); + constexpr int kOutputElements = + decltype(cute::size(output_fragment))::value; + static_assert( + (Traits::kTileN / Traits::kStoreTileN) * kOutputElements == + decltype(cute::size(accum0))::value, + "N32 store fragments must cover the native N256 accumulator"); + bool issuer = (tid_in_mma / Traits::kWarpSize) == 0; + int local_index = 0; + for (int index = cluster_idx; + index < total_cluster_tiles; + index += num_clusters, ++local_index) { + int stage = local_index % CommConfig::kNumStages; + int pass = local_index / CommConfig::kNumStages; + if (pass > 0) { + if (tid_in_mma == 0) { + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + cute::wait_barrier( + smem.dx_consumed_barrier[ + comm_warp][stage], + (pass - 1) & 1); + } + } + backward_consumer_barrier_sm90< + Config::kDxEpilogueBarrierId, Compute>(); + } + CUTE_UNROLL + for (int i = 0; i < decltype(cute::size(accum0))::value; ++i) { + asm volatile( + "mov.f32 %0, 0f00000000;" : "=f"(accum0(i))); + } + sm90::MainloopKLoopSm90::consume( + pipe, read_state, release_state, num_k_tiles, + [&](int stage) { + warpgroup_fence_operand(accum0); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, stage), + tCsB0(_, _, _, stage), accum0); + warpgroup_commit_batch(); + }, + [&]() { + warpgroup_fence_operand(accum0); + }); + + BackwardTileCoord tile = + backward_cluster_m_m_fast_tile( + index / DxSplitK, num_m_pairs, cluster_rank); + std::size_t slot_tile = + dx_slot_offset(cta, 0, stage) / + CommConfig::kTileElements; + + CUTE_UNROLL + for (int slice = 0; + slice < Traits::kTileN / Traits::kStoreTileN; + ++slice) { + if (slice > 0) { + if (issuer) { + wait_tma_store_complete(); + } + backward_consumer_barrier_sm90< + Config::kDxEpilogueBarrierId, Compute>(); + } + CUTE_UNROLL + for (int i = 0; i < kOutputElements; ++i) { + output_fragment(i) = + accum0(slice * kOutputElements + i); + } + copy(output_fragment, tDsD(_, _, _, Int<0>{})); + tma_store_fence(); + backward_consumer_barrier_sm90< + Config::kDxEpilogueBarrierId, Compute>(); + if (issuer && cute::elect_one_sync()) { + copy( + tma.dx_store, + tSsS(_, _, _, Int<0>{}), + tSgS( + _, _, _, + static_cast(slot_tile), + slice)); + tma_store_arrive(); + } + } + // Publish the completed FP32 tile before waking communication. + if (issuer) { + wait_tma_store_complete(); + liger_cute::detail::publish_local_reduce_source(); + } + backward_consumer_barrier_sm90< + Config::kDxEpilogueBarrierId, Compute>(); + if (tid_in_mma == 0) { + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + cute::arrive_barrier( + smem.dx_ready_epoch[comm_warp][stage]); + } + } + } + } + + int dw_n_tiles = Launch::num_dw_n_tiles(params.hidden); + int dw_n_pairs = ceil_div(dw_n_tiles, Traits::kClusterM); + int dw_total_units = + Launch::num_dw_m_tiles(params.local_vocab) * dw_n_pairs; + int dw_units = dw_total_units > cluster_idx + ? ceil_div(dw_total_units - cluster_idx, num_clusters) + : 0; + DwPhaseSm90:: + template consume_with_accumulator( + pipe, + read_state, + release_state, + accum0, + smem, + tma.dw_store, + tma.dw_add, + params, + wave, + dw_n_pairs, + kDwKTiles, + dw_units, + cluster_idx, + num_clusters, + cluster_rank); + } + + __syncthreads(); + sm90::cluster_exit_sm90(); +} + + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_sm90.cuh new file mode 100644 index 00000000..41a97eae --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_sm90.cuh @@ -0,0 +1,381 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy — backward GEMM contract. +// +// Structural port of +// src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_backward_fused_sm90.py +// (the CuTe-DSL `_FusedBackwardSM90` one-kernel backward). Constant names below +// map onto the reference module constants: +// +// THREADS_PER_CTA 384 tile_m 128 tile_n 256 tile_k 64 stages 3 +// m_tiles_per_wave 8 wave_rows 1024 +// dz_store_stages 1 -> 2 dw_store_stages 4 -> 2 +// epi_registers 88 -> 104 mma_registers 208 -> 200 +// +// The kernel keeps the reference's three fused phases and its wave loop: +// +// phase dZ Z = X @ W_local^T, softmax/entropy epilogue -> dZ wave buffer +// grid barrier +// phase dX dX_partial = dZ @ W_local, epilogue -> CTA-owned symmetric +// staging, the same CTA's warps 1..2 all-reduce and store dX +// grid barrier +// phase dW dW_local = dZ^T @ X, TMA store (wave 0) / reduce-add (wave > 0) +// grid barrier (unless this was the last wave) +// +// ── Deliberate deltas from the single-GPU reference ─────────────────────── +// * Warp roles are fixed by the tensor-parallel contract: warp 0 is the TMA +// producer for all three phases, warps 1..2 reduce that same CTA's dX groups, +// warp 3 is reserved and idle on every path, warps 4..11 are the two WGMMA +// consumer warp groups. The reference gave warp group 0 (warps 0..3) the dZ +// softmax epilogue; here warps 1..3 are unavailable, so the epilogue runs in +// the consumer warp groups straight out of the FP32 WGMMA accumulators. The +// reference's FP16 logit staging round-trip is reproduced in register +// (float -> half -> float) so the numerics are unchanged. +// * CLUSTER_M is 1, so the dX B operand is not multicast. The resident grid is +// launched through nvshmemx_collective_launch() for cross-PE co-scheduling; +// that API takes no cluster launch attribute. The reference's cluster-pair +// M multicast is therefore not expressible and is dropped. +// * The epilogue store tile is N64 rather than the reference's N32 for dZ/dW. +// A BF16 N64 row is exactly 128 B, which is what the SW128 store staging +// used by dZ and dW requires. dW already used N64 in the reference. +// * dX stays FP32 from WGMMA through a compact symmetric staging ring and the +// NVLS SUM. The communication warps convert once to X.dtype while +// scattering the final reduced result into the caller's row-major tensor +// (BF16 for this SM90 specialization). dW is untouched and stays rank-local. +// +// This header is deliberately CuTe-free: launch geometry, the wave/tile +// schedule, the warp plan and the launcher declarations only, so torch-free +// consumers can include the umbrella header without CUTLASS. The executable +// producer/consumer/mainloop lives in backward_gemm_mainloop_sm90.cuh and is +// instantiated by fused_linear_scaled_cross_entropy_backward.cu. +// ═══════════════════════════════════════════════════════════════════════════ + +#include + +#include +#include +#include + +#include "backward.cuh" +#include "config.cuh" +#include "dx_reduce.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +inline constexpr float kBackwardLog2E = 1.4426950408889634f; +inline constexpr int kBackwardMaxSmemBytes = 227 * 1024; + +// exp2 with the reference's fastmath=True lowering (llvm.nvvm.ex2.approx.ftz). +__host__ __device__ inline float backward_exp2_sm90(float value) { +#if defined(__CUDA_ARCH__) + float result; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); + return result; +#else + return exp2f(value); +#endif +} + +template +struct BackwardGemmConfigSm90 { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + static constexpr int kCompute = Compute; + + static constexpr int kTileM = 128; + static constexpr int kTileK = 64; + static constexpr int kDzTileN = 256; + static constexpr bool kDzNoTmaStore = false; + static constexpr int kDzLogicalTileN = kDzTileN; + static constexpr int kDxTileN = 256; + static constexpr int kDwTileN = 256; + static constexpr int kDwTileM = kTileM; + static constexpr int kDwLogicalTileN = kDwTileN; + static constexpr int kMainloopStages = 3; + static constexpr int kMTilesPerWave = 8; + static constexpr int kWaveRows = kTileM * kMTilesPerWave; + + static constexpr int kDzStoreTileN = 64; + static constexpr int kDzStoreStages = 2; + static constexpr int kDwStoreTileN = 64; + static constexpr int kDwStoreStages = 2; + + static constexpr int kClusterM = 1; + static constexpr int kNumThreads = + fused_scaled_linear_cross_entropy::kNumThreads; + static constexpr int kNumMmaWarpGroups = 2; + static constexpr int kWarpsPerWarpGroup = 4; + static constexpr int kWarpGroupSize = 128; + static constexpr int kConsumerThreads = + kNumMmaWarpGroups * kWarpGroupSize; + static constexpr int kFirstConsumerThread = kWarpGroupSize; + + // The dZ workspace K extent must be a whole number of K tiles so the dX / + // dW mainloops never see a ragged K tail on an operand they own. + static constexpr int kVocabAlign = kTileK; + + // setmaxnreg budget. Warps 0..3 share one warp group, so the communication + // warps take the producer allocation. The selected host-waved combined + // dX+dW specialization uses 56/224; separate phase-specialized kernels + // retain their budgets below. + static constexpr int kProducerRegisters = 56; + static constexpr int kMmaRegisters = 224; + static constexpr int kDzProducerRegisters = 24; + static constexpr int kDzMmaRegisters = 240; + static constexpr int kDxDwProducerRegisters = 72; + static constexpr int kDxDwMmaRegisters = 216; + static constexpr int kUsableRegisterBudget = 64512; + + // Hardware named barriers. 0 is __syncthreads(). + static constexpr int kDzEpilogueBarrierId = 4; // consumer warp groups + static constexpr int kDxEpilogueBarrierId = 5; // consumer warp groups + static constexpr int kDwStoreBarrierId = 6; // consumer warp groups + // Hands the aliased A/B and store-staging arenas from the dX phase to the + // dW phase. Deliberately NOT __syncthreads(): warps 1..2 are still + // draining the dX reduction ring and must not be pulled into it, and warp + // 3 is reserved. Only the producer warp and the two consumer warp groups + // touch the aliased storage, so only they take part. + static constexpr int kComputeBarrierId = 8; + static constexpr int kComputeBarrierThreads = + (1 + kNumMmaWarpGroups * kWarpsPerWarpGroup) * kWarpSize; + + static_assert(kNumThreads == (kNumMmaWarpGroups + 1) * kWarpGroupSize); + static_assert(kWaveRows % kTileM == 0); + static_assert(kDzTileN % kDzStoreTileN == 0); + static_assert(kComputeBarrierThreads == 288, + "the compute barrier covers warp 0 and warps 4..11"); + static_assert(kComputeBarrierThreads % kWarpSize == 0); + static_assert(kComputeBarrierThreads < kNumThreads, + "the compute barrier must exclude the communication and reserved " + "warps, otherwise it degenerates into __syncthreads()"); + static_assert(kDwTileN % kDwStoreTileN == 0); + static_assert(backward_warp_role(0) == BackwardWarpRole::kProducer); + static_assert(backward_warp_role(1) == BackwardWarpRole::kDxCommunication); + static_assert(backward_warp_role(2) == BackwardWarpRole::kDxCommunication); + static_assert(backward_warp_role(3) == BackwardWarpRole::kReserved); + static_assert(backward_warp_role(4) == BackwardWarpRole::kConsumer); + static_assert(backward_warp_role(11) == BackwardWarpRole::kConsumer); + + __host__ __device__ static constexpr int register_total() { + return kWarpGroupSize * kProducerRegisters + + kNumMmaWarpGroups * kWarpGroupSize * kMmaRegisters; + } + static_assert(register_total() <= kUsableRegisterBudget, + "SM90 backward register budget exceeds the usable per-CTA budget"); + static_assert( + kWarpGroupSize * kDzProducerRegisters + + kNumMmaWarpGroups * kWarpGroupSize * kDzMmaRegisters <= + kUsableRegisterBudget, + "SM90 dZ register budget exceeds the usable per-CTA budget"); + static_assert( + kWarpGroupSize * kDxDwProducerRegisters + + kNumMmaWarpGroups * kWarpGroupSize * kDxDwMmaRegisters <= + kUsableRegisterBudget, + "SM90 dX+dW register budget exceeds the usable per-CTA budget"); +}; + +// The reference's per-phase geometry, kept as named contracts so a future +// SM100 port can specialise them without touching the mainloop. +template +struct DzGemmContractSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + static constexpr int kTileM = Config::kTileM; + static constexpr int kTileN = Config::kDzTileN; + static constexpr int kTileK = Config::kTileK; + static constexpr int kStoreTileN = Config::kDzStoreTileN; + static constexpr int kStoreStages = Config::kDzStoreStages; +}; + +template +struct DxGemmContractSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + static constexpr int kTileM = Config::kTileM; + static constexpr int kTileN = Config::kDxTileN; + static constexpr int kTileK = Config::kTileK; +}; + +template +struct DwGemmContractSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + static constexpr int kTileM = Config::kTileM; + static constexpr int kTileN = Config::kDwTileN; + static constexpr int kTileK = Config::kTileK; + static constexpr int kStoreTileN = Config::kDwStoreTileN; + static constexpr int kStoreStages = Config::kDwStoreStages; +}; + +template +struct BackwardWarpPlanSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + static constexpr int kProducerWarp = 0; + static constexpr int kFirstDxCommWarp = 1; + static constexpr int kLastDxCommWarp = 2; + static constexpr int kReservedWarp = 3; + static constexpr int kFirstConsumerWarp = 4; + static constexpr int kLastConsumerWarp = 11; + // The consumer warp that issues every TMA store (the reference's `issuer`). + static constexpr int kIssuerWarp = kFirstConsumerWarp; + + static_assert(Config::kNumThreads == kNumThreads); + static_assert(kLastConsumerWarp + 1 == kNumWarps); + static_assert(backward_warp_role(kReservedWarp) == + BackwardWarpRole::kReserved); +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Problem description +// ─────────────────────────────────────────────────────────────────────────── + +// The vocabulary shard is contiguous: local column c is global vocabulary index +// `vocab_start + c`, and targets arrive as *global* int64 indices. `lse` and +// `entropy` are the forward's already-all-reduced per-token outputs, so they +// are ordinary local tensors — nothing symmetric is carried from forward. +template +struct BackwardGemmParamsSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + const void* x = nullptr; // BF16 [tokens, hidden] + const void* weight = nullptr; // BF16 [local_vocab, hidden] + const std::int64_t* target = nullptr; // int64 [tokens], global indices + const float* grad_output = nullptr; // FP32 [tokens], upstream NLL scale + const float* lse = nullptr; // FP32 [tokens], global LSE + const float* entropy = nullptr; // FP32 [tokens], global entropy + const float* entropy_grad = nullptr; // FP32 [tokens], upstream entropy scale + + void* grad_input = nullptr; // X.dtype [tokens, hidden], BF16 on SM90 + void* grad_weight = nullptr; // BF16 [local_vocab, hidden], rank local + + // Internal scratch, launcher owned (see BackwardSymmetricNames). + void* dz_workspace = nullptr; // BF16 [kWaveRows, padded_local_vocab] + std::size_t dz_workspace_bytes = 0; + + int tokens = 0; + int hidden = 0; + int local_vocab = 0; + std::int64_t vocab_start = 0; + std::int64_t ignore_index = -100; + float inverse_temperature = 1.0f; + // Selftest-only skew injection; production callers leave the defaults. + int dx_comm_delay_warp = -1; + int dx_comm_delay_iterations = 0; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Launch geometry +// ─────────────────────────────────────────────────────────────────────────── + +template +struct BackwardGemmLaunchSm90 { + using Config = BackwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + __host__ __device__ static constexpr int num_waves(int tokens) { + return ceil_div(tokens, Config::kWaveRows); + } + + // Vocabulary padded to a whole K tile so dX/dW never see a ragged K tail. + __host__ __device__ static constexpr int padded_vocab(int local_vocab) { + return ceil_div(local_vocab, Config::kVocabAlign) * Config::kVocabAlign; + } + + __host__ __device__ static constexpr int num_dz_n_tiles(int local_vocab) { + return ceil_div( + padded_vocab(local_vocab), Config::kDzLogicalTileN); + } + + __host__ __device__ static constexpr int num_dx_n_tiles(int hidden) { + return ceil_div(hidden, Config::kDxTileN); + } + + __host__ __device__ static constexpr int num_dx_k_tiles(int local_vocab) { + return padded_vocab(local_vocab) / Config::kTileK; + } + + __host__ __device__ static constexpr int num_dz_k_tiles(int hidden) { + return ceil_div(hidden, Config::kTileK); + } + + __host__ __device__ static constexpr int num_dw_m_tiles(int local_vocab) { + return ceil_div(padded_vocab(local_vocab), Config::kDwTileM); + } + + __host__ __device__ static constexpr int num_dw_n_tiles(int hidden) { + return ceil_div(hidden, Config::kDwLogicalTileN); + } + + __host__ __device__ static constexpr int num_dw_k_tiles() { + return Config::kWaveRows / Config::kTileK; + } + + // One wave of dZ, BF16. + __host__ __device__ static std::size_t dz_workspace_bytes( + int local_vocab) { + return static_cast(Config::kWaveRows) * + static_cast(padded_vocab(local_vocab)) * 2u; + } + + // Number of dX reduction groups produced by the whole launch, for a given + // TilesPerReduce. Waves are always full (`kMTilesPerWave` M tiles) so the + // group count is a pure function of the shape on every PE. + __host__ __device__ static int num_dx_groups( + int tokens, int hidden, int tiles_per_reduce) { + int groups_per_m = ceil_div(num_dx_n_tiles(hidden), tiles_per_reduce); + return num_waves(tokens) * Config::kMTilesPerWave * groups_per_m; + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Host launcher +// +// `params.grad_input` is filled with the TP-reduced dX and `params.grad_weight` +// with this rank's dW. Collective: every PE of the configured TP team calls +// this with the +// same shapes and TilesPerReduce. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct BackwardTpParamsSm90 { + static constexpr int kCompute = Compute; + + BackwardGemmParamsSm90 gemm = {}; + + // Compile-time knob selected at run time; 1, 2 (default) or 4. + int tiles_per_reduce = 2; + // Retained for FFI compatibility; CTA ownership ignores this value. + int num_comm_channels = 4; +}; + +__host__ __device__ constexpr int backward_dx_split_k( + int hidden, int local_vocab) { + // Split-K2 requires an even number of K64 vocabulary tiles. Ragged odd + // tile counts use split-K1 so the final tile is never dropped. + int k_tiles = (local_vocab + 63) / 64; + return hidden == 2048 && k_tiles % 2 == 0 ? 2 : 1; +} + +template +void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90& params, + cudaStream_t stream); + +extern template void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90<90>&, + cudaStream_t); +extern template void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90<90>&, + cudaStream_t); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu new file mode 100644 index 00000000..38d1eead --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu @@ -0,0 +1,753 @@ +#include "backward_reduce_sm90.cuh" + +#include +#include +#include + +#include +#include + +#include "backward_gemm_mainloop_sm90.cuh" +#include "liger_cute/check.h" +#include "liger_cute/detail/local_reduce.cuh" +#include "workspace.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { +namespace { + +using HostConfig = BackwardGemmConfigSm90<90>; +using HostLaunch = BackwardGemmLaunchSm90<90>; +using HostTraits = BackwardGemmTraitsSm90<90>; +using HostElement = HostTraits::Element; + +void check_cuda(cudaError_t error, const char* what) { + LIGER_CHECK( + error == cudaSuccess, + "fused_scaled_linear_cross_entropy backward reduction: ", + what, + " failed: ", + cudaGetErrorString(error)); +} + +template +auto make_row_major_load(const Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_LOAD{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_row_major_store(Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_STORE{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_row_major_reduce_add(Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_col_major_load( + const Element* ptr, int64_t rows, int64_t cols, int64_t leading) { + return make_tma_copy( + SM90_TMA_LOAD{}, + backward_col_major_tensor(ptr, rows, cols, leading), + Layout{}); +} + +template +__device__ void direct_peer_communicate( + BackwardDxDwSmemSm90& smem, + const BackwardGemmParamsSm90& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::DirectPeerReduceView& mapping, + int num_n_tiles, + int cta, + int num_ctas, + int wave, + int groups_per_wave) { + using Traits = BackwardGemmTraitsSm90; + using Config = typename Traits::Config; + + int comm_warp = + static_cast(threadIdx.x) / Traits::kWarpSize - + CommConfig::kFirstCommWarp; + int lane = static_cast(threadIdx.x) % Traits::kWarpSize; + auto* grad_input = static_cast<__nv_bfloat16*>(params.grad_input); + int next_index = 0; + + for (int unit = cta; unit < groups_per_wave; unit += num_ctas) { + int group_id = wave * groups_per_wave + unit; + DxTileGroup group = + dx_unit_to_group(unit, num_n_tiles, group_id); + DxCtaGroupSlot slot = + dx_cta_group_slot(next_index); + + std::uint64_t ready_epoch = + static_cast(group_id) + 1; + if (lane == 0) { + dx_ready_wait_acquire( + &smem.dx_ready_epoch[comm_warp][slot.stage], + ready_epoch); + } + __syncwarp(); + + std::size_t message_rows = + static_cast(group.num_tiles) * Traits::kTileM; + std::size_t first_row = + message_rows * static_cast(comm_warp) / + kDxCommWarpsPerChannel; + std::size_t last_row = + message_rows * static_cast(comm_warp + 1) / + kDxCommWarpsPerChannel; + std::size_t segment_elements = + (last_row - first_row) * Traits::kTileN; + std::size_t segment_begin = first_row * Traits::kTileN; + std::size_t base = + dx_slot_offset(cta, slot.comm_warp, slot.stage); + std::size_t ready_offset = dx_sync_offset( + cta, + comm_warp, + slot.stage, + kDxReadyPhase, + mapping.size); + std::size_t complete_offset = dx_sync_offset( + cta, + comm_warp, + slot.stage, + kDxCompletePhase, + mapping.size); + std::uint64_t epoch = + dx_epoch_base(comm) | + (static_cast(wave + 1) << 16) | + static_cast(slot.pass + 1); + + constexpr std::size_t kValuesPerVector = 4; + std::size_t segment_vectors = + segment_elements / kValuesPerVector; + float* local_reduced = comm.reduced + base + segment_begin; + liger_cute::detail::LocalReduceContext< + liger_cute::detail::LocalReduceBackend::kDirectPeer, + float> context{ + mapping.peer_partial, + base + segment_begin, + comm.sync + ready_offset, + mapping.peer_sync, + ready_offset, + comm.sync + complete_offset, + complete_offset, + mapping.rank, + mapping.size}; + liger_cute::detail::local_all_reduce< + liger_cute::detail::LocalReduceBackend::kDirectPeer>( + context, + local_reduced, + comm.partial + base + segment_begin, + segment_elements, + epoch); + + constexpr std::size_t kVectorsPerTile = + Traits::kTileM * Traits::kTileN / kValuesPerVector; + std::size_t segment_vector_begin = + segment_begin / kValuesPerVector; + for (std::size_t segment_vector = lane; + segment_vector < segment_vectors; + segment_vector += Traits::kWarpSize) { + std::size_t message_vector = + segment_vector_begin + segment_vector; + int tile = static_cast( + message_vector / kVectorsPerTile); + int vector_in_tile = static_cast( + message_vector % kVectorsPerTile); + int row = vector_in_tile / (Traits::kTileN / 4); + int column = + (vector_in_tile % (Traits::kTileN / 4)) * 4; + int output_row = wave * Config::kWaveRows + + group.m_tile * Traits::kTileM + row; + int output_column = + (group.first_n_tile + tile) * Traits::kTileN + column; + if (output_row < params.tokens && + output_column < params.hidden) { + float4 values = + reinterpret_cast( + local_reduced)[segment_vector]; + typename DxPhaseSm90< + CommConfig, Compute>::Bfloat16x4 packed{ + __floats2bfloat162_rn(values.x, values.y), + __floats2bfloat162_rn(values.z, values.w)}; + auto* output = reinterpret_cast< + typename DxPhaseSm90< + CommConfig, Compute>::Bfloat16x4*>( + grad_input + + static_cast(output_row) * + params.hidden + + output_column); + *output = packed; + } + } + __syncwarp(); + if (lane == 0) { + cute::arrive_barrier( + smem.dx_consumed_barrier[comm_warp][slot.stage]); + } + ++next_index; + } +} + +template < + bool RunDx, + bool RunDw, + bool ReturnEntropy, + int Compute, + class CommConfig, + class Bundle> +__global__ __launch_bounds__(BackwardGemmTraitsSm90::kNumThreads, 1) +void backward_dx_dw_direct_peer_wave_kernel_sm90( + __grid_constant__ const Bundle tma, + __grid_constant__ const BackwardGemmParamsSm90 params, + __grid_constant__ const DxReduceWorkspace comm, + __grid_constant__ const liger_cute::detail::DirectPeerReduceView mapping, + int wave) { + using Traits = BackwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Launch = BackwardGemmLaunchSm90; + using Smem = BackwardDxDwSmemSm90; + using PipelineTraits = typename Traits::PipelineTraits; + using ClusterShape = typename Traits::ClusterShape; + using Registers = sm90::WarpSpecializedRegistersSm90< + Config::kDxDwProducerRegisters, Config::kDxDwMmaRegisters, Compute>; + + extern __shared__ char raw_smem[]; + Smem& smem = *reinterpret_cast(raw_smem); + int warp_group = cutlass::canonical_warp_group_idx(); + int warp = cutlass::canonical_warp_idx_sync(); + + PipelineTraits::init_barriers( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + if (threadIdx.x == 0) { + CUTE_UNROLL + for (int comm_warp = 0; + comm_warp < kDxCommWarpsPerChannel; + ++comm_warp) { + CUTE_UNROLL + for (int stage = 0; stage < kDxRingStages; ++stage) { + smem.dx_ready_epoch[comm_warp][stage] = 0; + cute::initialize_barrier( + smem.dx_consumed_barrier[comm_warp][stage], 1); + } + } + } + sm90::cluster_pipeline_init_sm90(); + + int num_ctas = static_cast(gridDim.x); + int cta = static_cast(blockIdx.x); + if (warp_group == 0) { + Registers::producer(); + if (warp == 0) { + int padded_vocab = Launch::padded_vocab(params.local_vocab); + auto state = PipelineTraits::producer_start_state(); + auto pipe = PipelineTraits::make_producer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + threadIdx.x == 0, + ClusterShape{}); + if constexpr (RunDx) { + int dx_n_tiles = Launch::num_dx_n_tiles(params.hidden); + int dx_k_tiles = + Launch::num_dx_k_tiles(params.local_vocab); + int dx_groups_per_wave = Config::kMTilesPerWave * + dx_groups_per_m_tile(dx_n_tiles); + DxPhaseSm90:: + template produce( + pipe, + state, + smem, + tma.dz_load, + tma.wt, + params, + padded_vocab, + dx_n_tiles, + dx_k_tiles, + dx_groups_per_wave, + cta, + num_ctas); + } + + if constexpr (RunDx && RunDw) { + backward_compute_barrier_sm90(); + } + + if constexpr (RunDw) { + int dw_n_tiles = Launch::num_dw_n_tiles(params.hidden); + constexpr int kDwKTiles = Launch::num_dw_k_tiles(); + int dw_total = + Launch::num_dw_m_tiles(params.local_vocab) * dw_n_tiles; + int dw_tiles = dw_total > cta + ? ceil_div(dw_total - cta, num_ctas) + : 0; + DwPhaseSm90::template produce( + pipe, + state, + smem, + tma.dzt, + tma.xt, + params, + padded_vocab, + wave, + dw_n_tiles, + kDwKTiles, + dw_tiles, + cta, + num_ctas); + } + } else if ( + warp >= CommConfig::kFirstCommWarp && + warp <= CommConfig::kLastCommWarp) { + if constexpr (RunDx) { + int dx_n_tiles = Launch::num_dx_n_tiles(params.hidden); + int dx_groups_per_wave = Config::kMTilesPerWave * + dx_groups_per_m_tile(dx_n_tiles); + direct_peer_communicate( + smem, + params, + comm, + mapping, + dx_n_tiles, + cta, + num_ctas, + wave, + dx_groups_per_wave); + } + } + } else { + Registers::consumer(); + typename Traits::PipelineState read_state; + typename Traits::PipelineState release_state; + auto pipe = PipelineTraits::make_consumer( + smem.mainloop_pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + if constexpr (RunDx) { + int dx_n_tiles = Launch::num_dx_n_tiles(params.hidden); + int dx_k_tiles = + Launch::num_dx_k_tiles(params.local_vocab); + int dx_groups_per_wave = Config::kMTilesPerWave * + dx_groups_per_m_tile(dx_n_tiles); + int next_index = 0; + DxPhaseSm90:: + template consume( + pipe, + read_state, + release_state, + smem, + comm, + dx_n_tiles, + dx_k_tiles, + wave, + dx_groups_per_wave, + cta, + num_ctas, + next_index); + } + + if constexpr (RunDx && RunDw) { + backward_compute_barrier_sm90(); + } + + if constexpr (RunDw) { + int dw_n_tiles = Launch::num_dw_n_tiles(params.hidden); + constexpr int kDwKTiles = Launch::num_dw_k_tiles(); + int dw_total = + Launch::num_dw_m_tiles(params.local_vocab) * dw_n_tiles; + int dw_tiles = dw_total > cta + ? ceil_div(dw_total - cta, num_ctas) + : 0; + DwPhaseSm90::template consume( + pipe, + read_state, + release_state, + smem, + tma.dw_store, + tma.dw_add, + params, + wave, + dw_n_tiles, + kDwKTiles, + dw_tiles, + cta, + num_ctas); + } + } + + __syncthreads(); + sm90::cluster_exit_sm90(); +} + +template +__global__ void cluster_local_allgather_scatter_kernel( + __grid_constant__ const BackwardGemmParamsSm90 params, + __grid_constant__ const DxReduceWorkspace comm, + __grid_constant__ const liger_cute::detail::NvlsReduceView mapping, + int grid_ctas, + int wave, + int tiles_per_wave, + int num_n_tiles) { + using Traits = BackwardGemmTraitsSm90; + using Config = typename Traits::Config; + using CommConfig = DxCommConfig; + + int unit = static_cast(blockIdx.x); + if (unit >= tiles_per_wave) return; + int comm_warp = static_cast(threadIdx.x) / Traits::kWarpSize; + int lane = static_cast(threadIdx.x) % Traits::kWarpSize; + if (comm_warp >= kDxCommWarpsPerChannel) return; + + constexpr std::size_t kTileElements = CommConfig::kTileElements; + constexpr std::size_t kSegmentElements = + kTileElements / kDxCommWarpsPerChannel; + std::size_t packed_tile_elements = + kTileElements / mapping.size; + std::size_t packed_segment_elements = + kSegmentElements / mapping.size; + const float* packed_source = mapping.reduced_shard + + (static_cast(wave) * tiles_per_wave + unit) * + packed_tile_elements + + static_cast(comm_warp) * packed_segment_elements; + + int cta = unit % grid_ctas; + int next_index = unit / grid_ctas; + DxCtaGroupSlot slot = dx_cta_group_slot(next_index); + const float* values = packed_source; + if (mapping.size > 1) { + std::size_t ready_offset = dx_sync_offset( + cta, + comm_warp, + slot.stage, + kDxReadyPhase, + mapping.size); + std::size_t complete_offset = dx_sync_offset( + cta, + comm_warp, + slot.stage, + kDxCompletePhase, + mapping.size); + std::uint64_t epoch = + dx_epoch_base(comm) | 0x20000000u | + (static_cast(wave) << 16) | + static_cast(slot.pass + 1); + liger_cute::detail::nvls_barrier_warp( + comm.sync + ready_offset, + mapping.multicast_sync + ready_offset, + mapping.rank, + mapping.size, + epoch); + std::size_t full_tile_offset = + (static_cast(wave) * tiles_per_wave + unit) * + kTileElements; + std::size_t segment_begin = + static_cast(comm_warp) * kSegmentElements; + liger_cute::detail::nvls_allgather_warp( + mapping.multicast_reduced + + full_tile_offset + segment_begin, + packed_source, + kSegmentElements, + mapping.rank, + mapping.size); + liger_cute::detail::nvls_barrier_warp( + comm.sync + complete_offset, + mapping.multicast_sync + complete_offset, + mapping.rank, + mapping.size, + epoch); + values = comm.reduced + full_tile_offset + segment_begin; + } + + auto* grad_input = static_cast<__nv_bfloat16*>(params.grad_input); + std::size_t segment_begin = + static_cast(comm_warp) * kSegmentElements; + int m_tile = unit / num_n_tiles; + int n_tile = unit % num_n_tiles; + constexpr std::size_t kSegmentVectors = kSegmentElements / 4; + for (std::size_t vector = lane; + vector < kSegmentVectors; + vector += Traits::kWarpSize) { + std::size_t tile_vector = segment_begin / 4 + vector; + int row = static_cast( + tile_vector / (Traits::kTileN / 4)); + int column = static_cast( + (tile_vector % (Traits::kTileN / 4)) * 4); + int output_row = wave * Config::kWaveRows + + m_tile * Traits::kTileM + row; + int output_column = n_tile * Traits::kTileN + column; + if (output_row < params.tokens && + output_column < params.hidden) { + float4 value = + reinterpret_cast(values)[vector]; + typename DxPhaseSm90::Bfloat16x4 packed{ + __floats2bfloat162_rn(value.x, value.y), + __floats2bfloat162_rn(value.z, value.w)}; + auto* output = reinterpret_cast< + typename DxPhaseSm90< + CommConfig, Compute>::Bfloat16x4*>( + grad_input + + static_cast(output_row) * + params.hidden + + output_column); + *output = packed; + } + } +} + +template +void finalize_cluster_typed( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream) { + using CommConfig = DxCommConfig; + int num_n_tiles = HostLaunch::num_dx_n_tiles(params.hidden); + int tiles_per_wave = HostConfig::kMTilesPerWave * num_n_tiles; + std::size_t full_elements = + static_cast(num_waves) * + static_cast(tiles_per_wave) * + CommConfig::kTileElements; + std::size_t packed_elements = + full_elements / static_cast(mapping.size); + LIGER_CHECK( + full_elements * sizeof(float) <= + backward_dx_configured_durable_bytes(), + "fused_scaled_linear_cross_entropy backward: cluster dX durable " + "workspace capacity exceeded"); + if constexpr (RequiresRemote) { + liger_cute::detail::launch_remote_reduce( + remote, + comm.launch_epoch, + packed_elements, + stream); + } + + auto* kernel = &cluster_local_allgather_scatter_kernel<90>; + for (int wave = 0; wave < num_waves; ++wave) { + void* args[] = { + const_cast*>(¶ms), + const_cast*>(&comm), + const_cast(&mapping), + &grid, + &wave, + &tiles_per_wave, + &num_n_tiles}; + int status = nvshmemx_collective_launch( + reinterpret_cast(kernel), + dim3(static_cast(tiles_per_wave), 1, 1), + dim3(kDxCommWarpsPerChannel * kWarpSize, 1, 1), + args, + 0, + stream); + LIGER_CHECK( + status == 0, + "fused_scaled_linear_cross_entropy backward: cluster local " + "all-gather/scatter launch failed with status ", + status); + check_cuda( + cudaGetLastError(), + "cluster_local_allgather_scatter_kernel launch"); + } +} + +template < + int TilesPerReduce, + bool ReturnEntropy, + bool RunDx = true, + bool RunDw = true> +void launch_typed( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::DirectPeerReduceView& mapping, + int grid, + cudaStream_t stream, + int wave) { + using CommConfig = + DxCommConfig; + using Smem = BackwardDxDwSmemSm90<90>; + + auto* w = static_cast(params.weight); + auto* dz = static_cast(params.dz_workspace); + auto* dw = static_cast(params.grad_weight); + int padded_vocab = HostLaunch::padded_vocab(params.local_vocab); + + auto tma_dz_load = + make_row_major_load( + dz, HostConfig::kWaveRows, padded_vocab); + auto tma_wt = make_col_major_load< + typename HostTraits::SmemBMnMajor::Single>( + w, params.hidden, params.local_vocab, params.hidden); + auto tma_dzt = make_col_major_load< + typename HostTraits::SmemAMnMajor::Single>( + dz, padded_vocab, HostConfig::kWaveRows, padded_vocab); + auto tma_xt = make_col_major_load< + typename HostTraits::SmemBMnMajor::Single>( + static_cast(params.x), + params.hidden, + params.tokens, + params.hidden); + auto tma_dw_store = + make_row_major_store( + dw, params.local_vocab, params.hidden); + auto tma_dw_add = + make_row_major_reduce_add( + dw, params.local_vocab, params.hidden); + using Bundle = BackwardDxDwTmaBundleSm90< + decltype(tma_dz_load), + decltype(tma_wt), + decltype(tma_dzt), + decltype(tma_xt), + decltype(tma_dw_store), + decltype(tma_dw_add)>; + Bundle bundle{ + tma_dz_load, + tma_wt, + tma_dzt, + tma_xt, + tma_dw_store, + tma_dw_add}; + constexpr int kSmemBytes = static_cast(sizeof(Smem)); + static_assert( + 2 * kSmemBytes > kBackwardMaxSmemBytes, + "the direct-peer grid assumes at most one resident CTA per SM"); + auto* kernel = + &backward_dx_dw_direct_peer_wave_kernel_sm90< + RunDx, + RunDw, + ReturnEntropy, + 90, + CommConfig, + Bundle>; + check_cuda( + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + kSmemBytes), + "cudaFuncSetAttribute(MaxDynamicSharedMemorySize)"); + void* args[] = { + &bundle, + const_cast*>(¶ms), + const_cast*>(&comm), + const_cast(&mapping), + &wave}; + int status = nvshmemx_collective_launch( + reinterpret_cast(kernel), + dim3(static_cast(grid), 1, 1), + dim3(HostConfig::kNumThreads, 1, 1), + args, + static_cast(kSmemBytes), + stream); + LIGER_CHECK( + status == 0, + "fused_scaled_linear_cross_entropy backward: " + "non-NVLS collective launch failed with status ", + status); + check_cuda( + cudaGetLastError(), + "backward_dx_dw_direct_peer_wave_kernel_sm90 launch"); +} + +} // namespace + +void launch_backward_dx_dw_direct_peer_wave_sm90( + bool return_entropy, + int tiles_per_reduce, + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::DirectPeerReduceView& mapping, + int grid, + cudaStream_t stream, + int wave) { + LIGER_CHECK( + mapping.size > 1, + "fused_scaled_linear_cross_entropy backward: invalid direct-peer team " + "metadata"); + LIGER_CHECK( + mapping.available != 0, + "fused_scaled_linear_cross_entropy backward: NVLS is unavailable and " + "one or more tensor-parallel peers do not have a direct symmetric " + "mapping"); + if (return_entropy) { + switch (tiles_per_reduce) { + case 1: + launch_typed<1, true>( + params, comm, mapping, grid, stream, wave); + return; + case 2: + launch_typed<2, true>( + params, comm, mapping, grid, stream, wave); + return; + default: + launch_typed<4, true>( + params, comm, mapping, grid, stream, wave); + return; + } + } + switch (tiles_per_reduce) { + case 1: + launch_typed<1, false>( + params, comm, mapping, grid, stream, wave); + return; + case 2: + launch_typed<2, false>( + params, comm, mapping, grid, stream, wave); + return; + default: + launch_typed<4, false>( + params, comm, mapping, grid, stream, wave); + return; + } +} + +template +void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream) { + finalize_cluster_typed( + params, comm, mapping, remote, grid, num_waves, stream); +} + +template void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream); +template void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cuh new file mode 100644 index 00000000..6b5c67a8 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cuh @@ -0,0 +1,51 @@ +#pragma once + +// FSLCE-specific adapters around the reusable tensor-parallel reduction API. + +#include + +#include "backward_gemm_sm90.cuh" +#include "liger_cute/detail/tp_reduce.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +void launch_backward_dx_dw_direct_peer_wave_sm90( + bool return_entropy, + int tiles_per_reduce, + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::DirectPeerReduceView& mapping, + int grid, + cudaStream_t stream, + int wave); + +template +void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream); + +extern template void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream); +extern template void finalize_backward_dx_cluster_sm90( + const BackwardGemmParamsSm90<90>& params, + const DxReduceWorkspace& comm, + const liger_cute::detail::NvlsReduceView& mapping, + const liger_cute::detail::RemoteReduceView& remote, + int grid, + int num_waves, + cudaStream_t stream); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_sm90_selftest.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_sm90_selftest.cu new file mode 100644 index 00000000..470c38eb --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_sm90_selftest.cu @@ -0,0 +1,701 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// Internal self test for the tensor-parallel SM90 backward path. +// +// Compiled out unless LIGER_FSLCE_BACKWARD_SM90_SELFTEST is defined, so the +// core library's recursive .cu glob picks this file up as an empty translation +// unit. Build it with the production CMake objects so the non-RDC cluster/dW +// kernels and the isolated RDC fused fallback retain their production flags. +// NVSHMEM_REMOTE_TRANSPORT=none NVSHMEM_DISABLE_NCCL=1 \ +// NVSHMEM_SYMMETRIC_SIZE=3G \ +// nvshmrun -n 4 ./backward_selftest +// +// Diagnostics (environment variables, all optional): +// FSLCE_POISON=s|z|sz wipe the symmetric dX staging (s) and/or the dZ wave +// workspace (z) between the two attempts of a case, so +// a case that only passes on the second attempt cannot +// hide behind buffers it never wrote. +// FSLCE_CHANNELS= override the ignored legacy channel argument. +// FSLCE_DELAY_WARP=0|1 inject delay into one CTA-local comm warp. +// FSLCE_DELAY_ITERS=n number of 64-cycle sleeps before each owned chunk. +// FSLCE_ATTEMPTS=n repeat every case n times to expose ordering races. +// +// It runs at any PE count, including one: +// * dX is the FP32 TP *sum* over shards followed by one conversion to +// X.dtype (BF16 in this specialization), so every PE must end up with the same +// globally correct dX. The CPU reference rebuilds the whole vocabulary +// from the shared seed and sums every shard's contribution, so a dropped +// reduction, a wrong team, a mis-sliced tail group or a stale ring slot +// all show up. +// * dW is rank local, so it is checked against this shard's own tiles only. +// TilesPerReduce is swept, including one-tile chunks that alternate whole +// messages between the two comm warps. Legacy channel values are varied to +// verify that they no longer affect CTA ownership. +// ═══════════════════════════════════════════════════════════════════════════ + +#if defined(LIGER_FSLCE_BACKWARD_SM90_SELFTEST) + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "backward_gemm_sm90.cuh" +#include "buffer_pool.cuh" +#include "workspace.cuh" + +namespace { + +namespace fslce = liger::fused_scaled_linear_cross_entropy; + +using Element = __nv_bfloat16; + +int g_pe = 0; +int g_n_pes = 1; +int g_failures = 0; +constexpr int kDefaultAttempts = 2; +constexpr int kMaxTokens = 1280; +constexpr int kMaxHidden = 8448; +constexpr int kMaxLocalVocab = 1024; +constexpr int kMaxTilesPerReduce = 4; +constexpr int kLegacyMaxCommChannels = 4; + +void send_all(int fd, const void* data, std::size_t bytes) { + const char* cursor = static_cast(data); + while (bytes != 0) { + ssize_t sent = send(fd, cursor, bytes, 0); + if (sent <= 0) std::exit(1); + cursor += sent; + bytes -= static_cast(sent); + } +} + +void recv_all(int fd, void* data, std::size_t bytes) { + char* cursor = static_cast(data); + while (bytes != 0) { + ssize_t received = recv(fd, cursor, bytes, 0); + if (received <= 0) std::exit(1); + cursor += received; + bytes -= static_cast(received); + } +} + +void exchange_unique_id( + int rank, + int ranks, + const char* master, + int port, + nvshmemx_uniqueid_t* id) { + if (rank == 0) { + int server = socket(AF_INET, SOCK_STREAM, 0); + int reuse = 1; + setsockopt(server, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)); + sockaddr_in address = {}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_ANY); + address.sin_port = htons(static_cast(port)); + if (bind(server, reinterpret_cast(&address), sizeof(address)) != 0 || + listen(server, ranks) != 0) { + std::exit(1); + } + for (int remote = 1; remote < ranks; ++remote) { + int peer = accept(server, nullptr, nullptr); + if (peer < 0) std::exit(1); + send_all(peer, id, sizeof(*id)); + close(peer); + } + close(server); + return; + } + + sockaddr_in address = {}; + address.sin_family = AF_INET; + address.sin_port = htons(static_cast(port)); + if (inet_pton(AF_INET, master, &address.sin_addr) != 1) std::exit(1); + for (int attempt = 0; attempt < 200; ++attempt) { + int client = socket(AF_INET, SOCK_STREAM, 0); + if (connect( + client, + reinterpret_cast(&address), + sizeof(address)) == 0) { + recv_all(client, id, sizeof(*id)); + close(client); + return; + } + close(client); + usleep(100000); + } + std::exit(1); +} + +#define CUDA_OK(expr) \ + do { \ + cudaError_t status = (expr); \ + if (status != cudaSuccess) { \ + std::printf("[pe %d] CUDA error %s at %s:%d: %s\n", \ + g_pe, #expr, __FILE__, __LINE__, \ + cudaGetErrorString(status)); \ + std::exit(1); \ + } \ + } while (0) + +// Deterministic from (seed, index) so every PE reconstructs the same global +// tensors without any communication. +float element_at(std::uint64_t seed, std::size_t index, float scale) { + std::uint64_t state = seed ^ (index * 0x9E3779B97F4A7C15ULL); + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + state ^= state >> 29; + state = state * 0xBF58476D1CE4E5B9ULL; + state ^= state >> 32; + std::uint32_t bits = static_cast(state >> 33); + return (static_cast(bits % 2048) / 1024.0f - 1.0f) * scale; +} + +struct Reference { + std::vector dx; // [tokens, hidden], TP-summed + std::vector dw_local; // [local_vocab, hidden], this shard + std::vector lse; // [tokens] + std::vector entropy; // [tokens] +}; + +// Full-vocabulary CPU reference. Reproduces the kernel's FP16 logit staging +// round so the two agree on the softmax input bit for bit. +Reference reference( + const std::vector& x, + const std::vector& weight_global, + const std::vector& target, + const std::vector& grad_output, + const std::vector& entropy_grad, + int tokens, + int hidden, + int global_vocab, + int local_vocab, + int vocab_start, + std::int64_t ignore_index, + float inverse_temperature, + bool return_entropy) { + Reference out; + out.dx.assign(static_cast(tokens) * hidden, 0.0f); + out.dw_local.assign( + static_cast(local_vocab) * hidden, 0.0f); + out.lse.assign(tokens, 0.0f); + out.entropy.assign(tokens, 0.0f); + + std::vector scaled(global_vocab); + std::vector dz(global_vocab); + for (int m = 0; m < tokens; ++m) { + float row_max = -1.0e38f; + for (int v = 0; v < global_vocab; ++v) { + float dot = 0.0f; + for (int k = 0; k < hidden; ++k) { + dot += x[static_cast(m) * hidden + k] * + weight_global[static_cast(v) * hidden + k]; + } + // float -> half -> float, matching the reference's FP16 logit + // staging buffer. + scaled[v] = __half2float(__float2half(dot)) * inverse_temperature; + row_max = fmaxf(row_max, scaled[v]); + } + double sum = 0.0; + double weighted = 0.0; + for (int v = 0; v < global_vocab; ++v) { + double e = std::exp( + static_cast(scaled[v]) - static_cast(row_max)); + sum += e; + weighted += e * static_cast(scaled[v]); + } + float lse = row_max + static_cast(std::log(sum)); + out.lse[m] = lse; + out.entropy[m] = lse - static_cast(weighted / sum); + + bool ignored = target[m] == ignore_index; + float scale = ignored ? 0.0f : grad_output[m]; + float escale = (ignored || !return_entropy) ? 0.0f : entropy_grad[m]; + for (int v = 0; v < global_vocab; ++v) { + float value = 0.0f; + if (scale != 0.0f || escale != 0.0f) { + float probability = std::exp(scaled[v] - lse); + value = probability * + (scale + (lse - out.entropy[m] - scaled[v]) * escale); + } + if (!ignored && v == static_cast(target[m])) value -= scale; + value *= inverse_temperature; + // dZ is materialised in BF16 by the kernel. + dz[v] = __bfloat162float(__float2bfloat16(value)); + } + + for (int h = 0; h < hidden; ++h) { + float reduced = 0.0f; + for (int pe = 0; pe < global_vocab / local_vocab; ++pe) { + float partial = 0.0f; + int shard_start = pe * local_vocab; + for (int v = 0; v < local_vocab; ++v) { + int global_v = shard_start + v; + partial += dz[global_v] * + weight_global[ + static_cast(global_v) * hidden + h]; + } + reduced += partial; + } + out.dx[static_cast(m) * hidden + h] = reduced; + } + for (int v = 0; v < local_vocab; ++v) { + float value = dz[vocab_start + v]; + if (value == 0.0f) continue; + for (int h = 0; h < hidden; ++h) { + out.dw_local[static_cast(v) * hidden + h] += + value * x[static_cast(m) * hidden + h]; + } + } + } + return out; +} + +struct Comparison { + int bad; + float max_absolute; + float max_relative; +}; + +Comparison compare( + const char* label, + const std::vector& actual, + const std::vector& expected, + float absolute, + float relative, + int reported) { + Comparison result = {}; + for (std::size_t i = 0; i < expected.size(); ++i) { + float error = fabsf(actual[i] - expected[i]); + float relative_error = error / fmaxf(fabsf(expected[i]), 1.0e-6f); + result.max_absolute = fmaxf(result.max_absolute, error); + result.max_relative = fmaxf(result.max_relative, relative_error); + float tolerance = absolute + relative * fabsf(expected[i]); + if (!(error <= tolerance)) { + if (result.bad < reported) { + std::printf( + "[pe %d] FAIL %-4s [%zu]: got %12.6f want %12.6f " + "(tol %.6f)\n", + g_pe, label, i, actual[i], expected[i], tolerance); + } + ++result.bad; + } + } + return result; +} + +struct CaseShape { + int tokens; + int hidden; + int local_vocab; + float temperature; + std::uint64_t seed; +}; + +void check_chunk_ownership() { + using TestCommConfig = fslce::DxCommConfig< + fslce::BackwardGemmConfigSm90<90>, fslce::kDxRingStages, 2>; + constexpr int kGroupsPerWave = 136; + constexpr int kNumCtas = 132; + constexpr int kNumWaves = 3; + int seen[kNumCtas] = {}; + for (int group = 0; group < kNumWaves * kGroupsPerWave; ++group) { + int unit = group % kGroupsPerWave; + int cta = unit % kNumCtas; + int index = fslce::dx_cta_group_index( + group, kGroupsPerWave, kNumCtas); + fslce::DxCtaGroupSlot slot = + fslce::dx_cta_group_slot(index); + bool valid = index == seen[cta] && + slot.comm_warp == 0 && + slot.index_in_warp == index && + slot.stage == slot.index_in_warp % fslce::kDxRingStages && + slot.pass == slot.index_in_warp / fslce::kDxRingStages; + if (!valid) { + std::printf( + "[pe %d] invalid dX chunk owner: group=%d cta=%d " + "index=%d expected=%d warp=%d warp_index=%d stage=%d pass=%d\n", + g_pe, group, cta, index, seen[cta], slot.comm_warp, + slot.index_in_warp, slot.stage, slot.pass); + ++g_failures; + } + ++seen[cta]; + } +} + +template +void run_case( + const std::string& name, + const CaseShape& shape, + int tiles_per_reduce, + int num_comm_channels) { // NOLINT: mutated by FSLCE_CHANNELS + const char* filter = std::getenv("FSLCE_FILTER"); + if (filter != nullptr && name.find(filter) == std::string::npos) return; + + // ABI-compatibility diagnostic: this value no longer routes groups. + const char* channel_override = std::getenv("FSLCE_CHANNELS"); + if (channel_override != nullptr) { + num_comm_channels = std::atoi(channel_override); + } + const int tokens = shape.tokens; + const int hidden = shape.hidden; + const int local_vocab = shape.local_vocab; + const int global_vocab = local_vocab * g_n_pes; + const int vocab_start = g_pe * local_vocab; + const float scale = 3.0f / std::sqrt(static_cast(hidden)); + + std::vector host_x(static_cast(tokens) * hidden); + for (std::size_t i = 0; i < host_x.size(); ++i) { + host_x[i] = __bfloat162float( + __float2bfloat16(element_at(shape.seed, i, scale))); + } + std::vector host_w( + static_cast(global_vocab) * hidden); + for (std::size_t i = 0; i < host_w.size(); ++i) { + host_w[i] = __bfloat162float(__float2bfloat16( + element_at(shape.seed ^ 0xABCDEF01ULL, i, scale))); + } + + std::vector host_target(tokens); + std::vector host_grad_output(tokens); + std::vector host_entropy_grad(tokens); + for (int m = 0; m < tokens; ++m) { + if (m % 13 == 5) { + host_target[m] = -100; // ignore_index + } else { + std::uint64_t state = shape.seed ^ (0x5DEECE66DULL * (m + 1)); + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + host_target[m] = static_cast( + (state >> 17) % static_cast(global_vocab)); + } + host_grad_output[m] = 0.5f + 0.5f * element_at(shape.seed ^ 7u, m, 1.0f); + host_entropy_grad[m] = 0.25f * element_at(shape.seed ^ 11u, m, 1.0f); + } + + Reference expected = reference( + host_x, host_w, host_target, host_grad_output, host_entropy_grad, + tokens, hidden, global_vocab, local_vocab, vocab_start, -100, + 1.0f / shape.temperature, ReturnEntropy); + + // Device operands: X is replicated, W is this PE's contiguous shard. + std::vector device_x_host(host_x.size()); + for (std::size_t i = 0; i < host_x.size(); ++i) { + device_x_host[i] = __float2bfloat16(host_x[i]); + } + std::vector device_w_host( + static_cast(local_vocab) * hidden); + for (std::size_t i = 0; i < device_w_host.size(); ++i) { + device_w_host[i] = __float2bfloat16( + host_w[static_cast(vocab_start) * hidden + i]); + } + + Element* device_x = nullptr; + Element* device_w = nullptr; + Element* device_dx = nullptr; + Element* device_dw = nullptr; + std::int64_t* device_target = nullptr; + float* device_scalars = nullptr; // grad_output | lse | entropy | e_grad + CUDA_OK(cudaMalloc(&device_x, device_x_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_w, device_w_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_dx, device_x_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_dw, device_w_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_target, tokens * sizeof(std::int64_t))); + CUDA_OK(cudaMalloc(&device_scalars, 4 * tokens * sizeof(float))); + CUDA_OK(cudaMemcpy(device_x, device_x_host.data(), + device_x_host.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_w, device_w_host.data(), + device_w_host.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_target, host_target.data(), + tokens * sizeof(std::int64_t), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_scalars + 0 * tokens, host_grad_output.data(), + tokens * sizeof(float), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_scalars + 1 * tokens, expected.lse.data(), + tokens * sizeof(float), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_scalars + 2 * tokens, expected.entropy.data(), + tokens * sizeof(float), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_scalars + 3 * tokens, host_entropy_grad.data(), + tokens * sizeof(float), cudaMemcpyHostToDevice)); + // Poison the outputs so a skipped store is a failure, not a pass. + CUDA_OK(cudaMemset(device_dx, 0x7f, + device_x_host.size() * sizeof(Element))); + CUDA_OK(cudaMemset(device_dw, 0x7f, + device_w_host.size() * sizeof(Element))); + + fslce::BackwardScratch scratch = + fslce::reserve_backward_scratch(local_vocab); + + // Diagnostic: FSLCE_POISON=s wipes the symmetric dX staging between + // attempts, =z wipes the dZ wave workspace, =sz both. Any case that only + // passes without poisoning is reading a buffer it never wrote. + const char* poison = std::getenv("FSLCE_POISON"); + bool poison_staging = poison != nullptr && std::strchr(poison, 's'); + bool poison_dz = poison != nullptr && std::strchr(poison, 'z'); + fslce::DxReduceWorkspace staging = + fslce::reserve_dx_reduce_workspace( + tiles_per_reduce, + fslce::kDxRingStages, + fslce::backward_dx_resident_cta_capacity()); + std::size_t staging_bytes = + fslce::backward_dx_staging_bytes(kMaxTilesPerReduce); + + fslce::BackwardTpParamsSm90<90> params; + params.gemm.x = device_x; + params.gemm.weight = device_w; + params.gemm.target = device_target; + params.gemm.grad_output = device_scalars + 0 * tokens; + params.gemm.lse = device_scalars + 1 * tokens; + params.gemm.entropy = device_scalars + 2 * tokens; + params.gemm.entropy_grad = device_scalars + 3 * tokens; + params.gemm.grad_input = device_dx; + params.gemm.grad_weight = device_dw; + params.gemm.dz_workspace = scratch.dz_workspace; + params.gemm.dz_workspace_bytes = scratch.dz_workspace_bytes; + params.gemm.tokens = tokens; + params.gemm.hidden = hidden; + params.gemm.local_vocab = local_vocab; + params.gemm.vocab_start = vocab_start; + params.gemm.ignore_index = -100; + params.gemm.inverse_temperature = 1.0f / shape.temperature; + const char* delay_warp = std::getenv("FSLCE_DELAY_WARP"); + const char* delay_iters = std::getenv("FSLCE_DELAY_ITERS"); + const char* attempts_env = std::getenv("FSLCE_ATTEMPTS"); + int attempts = attempts_env == nullptr + ? kDefaultAttempts + : std::atoi(attempts_env); + if (attempts < 1) attempts = 1; + if (delay_warp != nullptr) { + params.gemm.dx_comm_delay_warp = std::atoi(delay_warp); + } + if (delay_iters != nullptr) { + params.gemm.dx_comm_delay_iterations = std::atoi(delay_iters); + } + params.tiles_per_reduce = tiles_per_reduce; + params.num_comm_channels = num_comm_channels; + + // Each case is launched twice: a first-launch-only failure is a warm-up + // bug, a failure that moves between repetitions is a race. + cudaEvent_t start; + cudaEvent_t stop; + CUDA_OK(cudaEventCreate(&start)); + CUDA_OK(cudaEventCreate(&stop)); + for (int attempt = 0; attempt < attempts; ++attempt) { + if (poison_staging) { + CUDA_OK(cudaMemset(staging.partial, 0, staging_bytes)); + CUDA_OK(cudaMemset(staging.reduced, 0, staging_bytes)); + } + if (poison_dz) { + CUDA_OK(cudaMemset( + scratch.dz_workspace, 0, scratch.dz_workspace_bytes)); + } + CUDA_OK(cudaMemset(device_dx, 0x7f, + device_x_host.size() * sizeof(Element))); + CUDA_OK(cudaMemset(device_dw, 0x7f, + device_w_host.size() * sizeof(Element))); + nvshmem_barrier_all(); + CUDA_OK(cudaEventRecord(start)); + fslce::fused_linear_scaled_cross_entropy_backward< + ReturnEntropy, 90>(params, nullptr); + CUDA_OK(cudaEventRecord(stop)); + CUDA_OK(cudaGetLastError()); + CUDA_OK(cudaEventSynchronize(stop)); + float elapsed_ms = 0.0f; + CUDA_OK(cudaEventElapsedTime(&elapsed_ms, start, stop)); + nvshmem_barrier_all(); + + std::vector raw_dx(device_x_host.size()); + std::vector raw_dw(device_w_host.size()); + CUDA_OK(cudaMemcpy(raw_dx.data(), device_dx, + raw_dx.size() * sizeof(Element), cudaMemcpyDeviceToHost)); + CUDA_OK(cudaMemcpy(raw_dw.data(), device_dw, + raw_dw.size() * sizeof(Element), cudaMemcpyDeviceToHost)); + std::vector got_dx(raw_dx.size()); + std::vector got_dw(raw_dw.size()); + for (std::size_t i = 0; i < raw_dx.size(); ++i) { + got_dx[i] = __bfloat162float(raw_dx[i]); + } + for (std::size_t i = 0; i < raw_dw.size(); ++i) { + got_dw[i] = __bfloat162float(raw_dw[i]); + } + + Comparison dx = compare( + "dX", got_dx, expected.dx, 6.0e-3f, 4.0e-2f, 4); + Comparison dw = compare( + "dW", got_dw, expected.dw_local, 8.0e-3f, 4.0e-2f, 4); + g_failures += dx.bad + dw.bad; + + if (g_pe == 0) { + std::printf( + "%-42s #%d %s (PEs=%d M=%d H=%d Vloc=%d T=%.2f entropy=%d " + "TPR=%d legacy_chan=%d dx_bad=%d dw_bad=%d ms=%.3f " + "dx_max_abs=%.6g dx_max_rel=%.6g " + "dw_max_abs=%.6g dw_max_rel=%.6g)\n", + name.c_str(), attempt, + (dx.bad + dw.bad) == 0 ? "PASS" : "FAIL", + g_n_pes, tokens, hidden, local_vocab, shape.temperature, + static_cast(ReturnEntropy), tiles_per_reduce, + num_comm_channels, dx.bad, dw.bad, elapsed_ms, + dx.max_absolute, dx.max_relative, + dw.max_absolute, dw.max_relative); + } + } + CUDA_OK(cudaEventDestroy(start)); + CUDA_OK(cudaEventDestroy(stop)); + + CUDA_OK(cudaFree(device_x)); + CUDA_OK(cudaFree(device_w)); + CUDA_OK(cudaFree(device_dx)); + CUDA_OK(cudaFree(device_dw)); + CUDA_OK(cudaFree(device_target)); + CUDA_OK(cudaFree(device_scalars)); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc == 6) { + int rank = std::atoi(argv[1]); + int ranks = std::atoi(argv[2]); + int device = std::atoi(argv[3]); + CUDA_OK(cudaSetDevice(device)); + CUDA_OK(cudaFree(nullptr)); + nvshmemx_uniqueid_t id = NVSHMEMX_UNIQUEID_INITIALIZER; + if (rank == 0 && nvshmemx_get_uniqueid(&id) != 0) return 1; + exchange_unique_id( + rank, ranks, argv[4], std::atoi(argv[5]), &id); + nvshmemx_init_attr_t attributes = NVSHMEMX_INIT_ATTR_INITIALIZER; + nvshmemx_set_attr_uniqueid_args( + rank, ranks, &id, &attributes); + if (nvshmemx_init_attr( + NVSHMEMX_INIT_WITH_UNIQUEID, &attributes) != 0) { + return 1; + } + } else { + nvshmem_init(); + } + g_pe = nvshmem_my_pe(); + g_n_pes = nvshmem_n_pes(); + CUDA_OK(cudaSetDevice(nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE))); + + cudaDeviceProp properties = {}; + CUDA_OK(cudaGetDeviceProperties( + &properties, nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE))); + if (properties.major != 9) { + if (g_pe == 0) { + std::printf("skipping: requires Hopper (sm90), found sm_%d%d\n", + properties.major, properties.minor); + } + nvshmem_finalize(); + return 0; + } + + if (g_pe == 0) { + std::printf("tensor-parallel backward over %d PEs\n", g_n_pes); + } + check_chunk_ownership(); + + // Capacities and multicast mappings are immutable once allocated, so + // configure for the largest shape and the widest knobs used below. + // Collective. + fslce::configure_backward_tp_symmetric( + kMaxTokens, kMaxHidden, kMaxLocalVocab, + kMaxTilesPerReduce, kLegacyMaxCommChannels, + static_cast(NVSHMEM_TEAM_WORLD)); + if (g_pe == 0) { + std::printf( + "symmetric staging %.2f MiB, device scratch %.2f MiB, " + "resident_ctas=%d nvls_team_size=%d\n", + fslce::backward_tp_pool_symmetric_bytes( + kMaxTokens, kMaxHidden, + kMaxTilesPerReduce, kLegacyMaxCommChannels) / + (1024.0 * 1024.0), + fslce::backward_tp_pool_device_bytes(kMaxLocalVocab) / + (1024.0 * 1024.0), + fslce::backward_dx_resident_cta_capacity(), + fslce::backward_dx_team_size()); + } + + // One N256 hidden tile: no grouping, exercises the simplest ring path. + const CaseShape small{256, 256, 512, 1.0f, 101}; + // Three N256 hidden tiles with a ragged last one (648 = 2 * 256 + 136), so + // TilesPerReduce 2 produces a {2, 1} tail split and 4 a {3} tail. + const CaseShape ragged{200, 648, 333, 0.8f, 103}; + // Two waves: exercises the dW TMA reduce-add path and a ring that wraps. + const CaseShape two_waves{1280, 512, 512, 1.0f, 105}; + // Five K64 vocabulary tiles: H2048 must fall back to split-K1 rather than + // dropping the final odd tile. Non-unit temperature checks the chain rule. + const CaseShape odd_k_tiles{128, 2048, 320, 0.8f, 106}; + // 136 dX groups with TPR=1 on a 132-SM H100: launches the full resident + // grid and makes four CTAs wrap to a second CTA-owned group. + const CaseShape full_grid{1, 4352, 64, 1.0f, 107}; + // 136 two-tile chunks: the same full-grid wrap, now proving adjacent + // grouped messages alternate between the CTA's two comm warps. + const CaseShape full_grid_grouped{1, 8448, 64, 1.0f, 109}; + + try { + run_case("dX/dW aligned, no entropy", small, 2, 2); + run_case("dX/dW aligned, entropy", small, 2, 2); + run_case("dX/dW ragged M/H/V, TilesPerReduce=1", ragged, 1, 2); + run_case("dX/dW ragged M/H/V, TilesPerReduce=2", ragged, 2, 2); + run_case("dX/dW ragged M/H/V, TilesPerReduce=4", ragged, 4, 2); + run_case("dX/dW ragged, legacy channel arg=1", ragged, 2, 1); + run_case("dX/dW ragged, legacy channel arg=4", ragged, 2, 4); + run_case("dX/dW two waves, TilesPerReduce=2", two_waves, 2, 4); + run_case( + "dX/dW H2048 odd K tiles and temperature", + odd_k_tiles, + 1, + 2); + run_case("dX/dW full resident CTA grid", full_grid, 1, 2); + run_case( + "dX/dW full grid, alternating two-tile chunks", + full_grid_grouped, 2, 2); + } catch (const std::exception& error) { + std::printf("[pe %d] exception: %s\n", g_pe, error.what()); + fslce::reset_fslce_tp_configuration(); + liger::global_buffer_pool().clear(); + nvshmem_finalize(); + return 1; + } + + // Make every rank's verdict visible before PE 0 prints the summary. + // Pooled, not raw nvshmem_malloc: this module never runs its own symmetric + // allocation lifecycle. + auto& pool = liger::global_buffer_pool(); + int* local_slot = static_cast( + pool.get_symmetric("fslce_bwd_selftest_verdict_local", sizeof(int))); + int* total_slot = static_cast( + pool.get_symmetric("fslce_bwd_selftest_verdict_total", sizeof(int))); + CUDA_OK(cudaMemcpy( + local_slot, &g_failures, sizeof(int), cudaMemcpyHostToDevice)); + nvshmem_int_sum_reduce(NVSHMEM_TEAM_WORLD, total_slot, local_slot, 1); + int host_total = 0; + CUDA_OK(cudaMemcpy( + &host_total, total_slot, sizeof(int), cudaMemcpyDeviceToHost)); + + if (g_pe == 0) { + std::printf("%s: %d comparison failures across %d PEs\n", + host_total == 0 ? "OK" : "FAILED", host_total, g_n_pes); + } + // Release the pooled allocations while NVSHMEM is still up: the pool is a + // function-local static whose destructor would otherwise run after + // nvshmem_finalize(). + fslce::reset_fslce_tp_configuration(); + liger::global_buffer_pool().clear(); + nvshmem_finalize(); + return host_total == 0 ? 0 : 1; +} + +#endif // LIGER_FSLCE_BACKWARD_SM90_SELFTEST diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/config.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/config.cuh new file mode 100644 index 00000000..411bc09e --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/config.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +inline constexpr int kWarpSize = 32; +inline constexpr int kNumThreads = 384; +inline constexpr int kNumWarps = kNumThreads / kWarpSize; + +enum class BackwardWarpRole : std::uint8_t { + kProducer, + kDxCommunication, + kReserved, + kConsumer, + kInactive, +}; + +__host__ __device__ constexpr BackwardWarpRole backward_warp_role(int warp_id) { + return warp_id == 0 + ? BackwardWarpRole::kProducer + : (warp_id >= 1 && warp_id <= 2 + ? BackwardWarpRole::kDxCommunication + : (warp_id == 3 + ? BackwardWarpRole::kReserved + : (warp_id >= 4 && warp_id < kNumWarps + ? BackwardWarpRole::kConsumer + : BackwardWarpRole::kInactive))); +} + +struct DxTileGroup { + int m_tile; + int first_n_tile; + int num_tiles; + int sequence; +}; + +__host__ __device__ constexpr int ceil_div(int value, int divisor) { + return (value + divisor - 1) / divisor; +} + +template +__host__ __device__ constexpr int dx_groups_per_m_tile(int num_n_tiles) { + return ceil_div(num_n_tiles, Config::kCoalesceTiles); +} + +template +__host__ __device__ constexpr DxTileGroup make_dx_tile_group( + int m_tile, int group_in_m, int num_n_tiles, int sequence) { + int first_n_tile = group_in_m * Config::kCoalesceTiles; + int remaining = num_n_tiles - first_n_tile; + int num_tiles = remaining < Config::kCoalesceTiles + ? remaining + : Config::kCoalesceTiles; + return {m_tile, first_n_tile, num_tiles > 0 ? num_tiles : 0, sequence}; +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/dx_reduce.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/dx_reduce.cuh new file mode 100644 index 00000000..8443bd8f --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/dx_reduce.cuh @@ -0,0 +1,240 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// dX tensor-parallel reduction contract — CTA-owned staging, tile grouping +// and one-team NVLS synchronization. +// +// The fused SM90 backward kernel computes a *partial* dX on every rank: the +// vocabulary is sharded, so `dX_local = dZ_local @ W_local` is one term of a +// SUM over the tensor-parallel team. Instead of materialising the whole +// partial dX and running one big collective afterwards, each completed +// M128xN256 FP32 accumulator tile is written straight into a compact symmetric +// FP32 staging ring owned by that same CTA. Its own warps 1..2 reduce and +// scatter the group while the WGMMA consumers keep working on later groups. +// +// ── Pipeline ────────────────────────────────────────────────────────────── +// consumer epilogue FP32 accumulator -> generic coalesced global stores +// into partial[slot], CTA named-barrier release, +// then one CTA-scope release-store of a unique epoch +// owning comm warp local_chunk_index % 2 selects warp 1 or 2; +// it acquire-loads that epoch, publishes a block-prefixed +// system-scope NVLS ready epoch, performs the two-shot +// multicast SUM, then publishes a completion epoch +// owning comm warp vector-load the complete FP32 result, +// convert once to X.dtype (BF16 in this specialization), +// scatter into row-major dX, +// CTA-local consumed mbarrier arrival +// +// ── Grouping ────────────────────────────────────────────────────────────── +// `TilesPerReduce` contiguous N tiles of one M tile form one reduction +// message. The last group of an M tile is a ragged tail of +// `num_n_tiles % TilesPerReduce` tiles; its message is exactly that many tiles +// long, which is identical on every PE because every PE sees the same shapes. +// +// ── CTA ownership and independent NVLS regions ──────────────────────────── +// Wave-local group `unit` belongs to CTA `unit % gridDim.x`; warp 0 and warps +// 4..11 produce only that CTA's groups, and warps 1..2 consume exactly the +// same sequence. Ring slot `j % NumStages` is therefore private to one CTA, +// and ready/reuse handoff is entirely CTA-local. +// +// Every ring slot has disjoint symmetric data and signal prefixes. All CTAs +// share one tensor-parallel team's immutable NVLS multicast mapping; no +// collective counter or team scratch is used inside the fused kernel. +// +// Warp 3 is reserved and idle on every path. +// +// This header stays NVSHMEM-free and CUTLASS-free so torch-free consumers can +// include the umbrella header. +// ═══════════════════════════════════════════════════════════════════════════ + +#include + +#include +#include + +#include "config.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +// Hopper has at most 132 resident SMs/CTAs for this 1-CTA/SM kernel. +inline constexpr int kMaxDxResidentCtas = 132; +inline constexpr int kDxCommWarpsPerChannel = 2; +inline constexpr int kDxSyncPhases = 2; +inline constexpr int kDxReadyPhase = 0; +inline constexpr int kDxCompletePhase = 1; + +// Ring depth. Compile time: the kernel template, symmetric capacity and +// CTA-local mbarrier arrays all have to agree, so there is exactly one place +// to change it. 2, 3 and 4 are supported and exercised; 4 is the default +// because the comm warps run one pipeline step behind the collective, so a +// producer can run ahead by kDxRingStages - 1 groups. +inline constexpr int kDxRingStages = 4; +static_assert(kDxRingStages >= 2, + "the dX staging ring needs at least a double buffer"); + +// Compile-time shape of the staging ring. +// +// NumStages ring depth per CTA; >= 2 or the producer of a group +// would wait on the very slot it is about to fill. +// TilesPerReduce contiguous N tiles coalesced into one reduction message. +template < + typename GemmConfig, + int NumStages = 4, + int TilesPerReduce = 2, + int Compute = GemmConfig::kCompute> +struct DxCommConfig { + static_assert(Compute == 90, "DxCommConfig requires Compute=90"); + static_assert(NumStages >= 2, + "the dX staging ring needs at least a double buffer; with one stage " + "the producer of group j would wait on group j itself"); + static_assert(TilesPerReduce >= 1, + "TilesPerReduce must be positive"); + + static constexpr int kCompute = Compute; + static constexpr int kNumStages = NumStages; + static constexpr int kTilesPerReduce = TilesPerReduce; + // config.cuh's make_dx_tile_group()/dx_groups_per_m_tile() spell this + // knob kCoalesceTiles; keep the alias so both names name one value. + static constexpr int kCoalesceTiles = TilesPerReduce; + + static constexpr int kTileM = GemmConfig::kTileM; + static constexpr int kTileN = GemmConfig::kDxTileN; + static constexpr int kTileElements = kTileM * kTileN; + static constexpr int kGroupElements = kTileElements * kTilesPerReduce; + + static constexpr int kProducerWarp = 0; + static constexpr int kFirstCommWarp = 1; + static constexpr int kNumCommWarps = kDxCommWarpsPerChannel; + static constexpr int kLastCommWarp = kFirstCommWarp + kNumCommWarps - 1; + static constexpr int kReservedWarp = 3; + + static_assert(kLastCommWarp + 1 == kReservedWarp, + "warps 1..2 cooperatively communicate and warp 3 stays reserved"); + static_assert(backward_warp_role(kProducerWarp) == + BackwardWarpRole::kProducer); + static_assert(backward_warp_role(kFirstCommWarp) == + BackwardWarpRole::kDxCommunication); + static_assert(backward_warp_role(kLastCommWarp) == + BackwardWarpRole::kDxCommunication); + static_assert(backward_warp_role(kReservedWarp) == + BackwardWarpRole::kReserved); +}; + +// FSLCE-specific symmetric ring storage. Backend mappings are supplied through +// the minimal views in liger_cute/detail/tp_reduce.cuh. +// +// partial / reduced are symmetric (NVSHMEM) and laid out identically on every +// PE as +// [max_resident_ctas][2 comm warps][kNumStages][kTilesPerReduce][kTileM][kTileN], +// FP32 and +// fully packed, so one group is one contiguous message. FP32 is intentional: +// TP=1 and TP>1 use the same reduction path and accumulation semantics, and +// the final row-major grad_input store is the only conversion to X.dtype +// (BF16 in the SM90 specialization). +// +template +struct DxReduceWorkspace { + Element* partial; // symmetric, SUM source + Element* reduced; // symmetric, SUM destination + std::uint64_t* sync; // local symmetric signal replicas + const std::uint64_t* launch_epoch; // device-resident graph-safe generation +}; + +#if defined(__CUDACC__) +template +__device__ __forceinline__ std::uint64_t dx_epoch_base( + const DxReduceWorkspace& workspace) { + return *workspace.launch_epoch; +} +#endif + +template +__host__ __device__ inline std::size_t dx_slot_index( + int cta, int comm_warp, int stage) { + return (static_cast(cta) * kDxCommWarpsPerChannel + + comm_warp) * + Config::kNumStages + + stage; +} + +// Element offset of a slot's first tile inside the packed staging arena. The +// compile-time communication config is the single source of truth for the +// staging stride used by all producers and consumers. +template +__host__ __device__ inline std::size_t dx_slot_offset( + int cta, int comm_warp, int stage) { + std::size_t slot = dx_slot_index(cta, comm_warp, stage); + return slot * static_cast(Config::kTilesPerReduce) * + static_cast(Config::kTileElements); +} + +template +__host__ __device__ inline std::size_t dx_sync_offset( + int cta, + int comm_warp, + int stage, + int phase, + int team_size) { + std::size_t slot = dx_slot_index(cta, comm_warp, stage); + return (slot * kDxSyncPhases + static_cast(phase)) * + static_cast(team_size); +} + +// ─────────────────────────────────────────────────────────────────────────── +// Deterministic group schedule +// ─────────────────────────────────────────────────────────────────────────── + +// Index in one CTA's deterministic group sequence -> ring coordinates. +// Both communication warps cooperatively reduce disjoint halves of the same +// slot, so the payload always lives in data slot zero. +struct DxCtaGroupSlot { + int comm_warp; + int index_in_warp; + int stage; + int pass; +}; + +template +__host__ __device__ inline DxCtaGroupSlot dx_cta_group_slot( + int index_in_cta) { + DxCtaGroupSlot slot; + slot.comm_warp = 0; + slot.index_in_warp = index_in_cta; + slot.stage = slot.index_in_warp % Config::kNumStages; + slot.pass = slot.index_in_warp / Config::kNumStages; + return slot; +} + +// Wave-local unit index -> the (m_tile, first_n_tile, num_tiles) it covers. +// `unit` counts M-fastest-last: unit = m_tile * groups_per_m + group_in_m. +template +__host__ __device__ inline DxTileGroup dx_unit_to_group( + int unit, int num_n_tiles, int group_id) { + int groups_per_m = dx_groups_per_m_tile(num_n_tiles); + int m_tile = unit / groups_per_m; + int group_in_m = unit - m_tile * groups_per_m; + return make_dx_tile_group( + m_tile, group_in_m, num_n_tiles, group_id); +} + +__host__ __device__ inline int dx_groups_for_cta_per_wave( + int groups_per_wave, int num_ctas, int cta) { + return groups_per_wave > cta + ? (groups_per_wave - cta + num_ctas - 1) / num_ctas + : 0; +} + +// Global group -> index in its owning CTA's sequence across all waves. +__host__ __device__ inline int dx_cta_group_index( + int group_id, int groups_per_wave, int num_ctas) { + int wave = group_id / groups_per_wave; + int unit = group_id - wave * groups_per_wave; + int cta = unit % num_ctas; + int index_in_wave = unit / num_ctas; + return wave * dx_groups_for_cta_per_wave( + groups_per_wave, num_ctas, cta) + index_in_wave; +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_mainloop_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_mainloop_sm90.cuh new file mode 100644 index 00000000..2107f55a --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_mainloop_sm90.cuh @@ -0,0 +1,794 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy — executable forward mainloop. +// +// Structural port of _fused_scaled_cross_entropy_forward_fragment_sm90.py. +// The CuTe/CUTLASS mechanics (TiledMMA construction, PipelineTmaAsync driving, +// TMA partitioning, warpgroup fences) follow the repository's SM90 idioms in +// csrc/core/src/moe, but the loop nest, work assignment, staging order, +// masking and online-softmax sequencing are a literal transcription of the +// reference kernel — not a generic CUTLASS GEMM. +// +// CTA layout (384 threads = 12 warps, cluster = (2, 1, 1)): +// warp 0 : TMA producer, 24 registers. Loads the CTA's own M128xK64 X +// tile and multicasts two N160xK64 W panels to the cluster +// peer through one 3-stage mbarrier pipeline. +// warps 1..3 : idle in forward (they carry the dX comms in backward), but +// they still take PRODUCER_REGISTERS so the warp group's +// setmaxnreg is uniform. +// warps 4..11 : two WGMMA consumer warp groups, 240 registers. Each owns 64 +// token rows and two FP32 N160 accumulators. +// +// Mainloop: the consumer read state and the release state are deliberately +// separate. Stage zero is issued without a release; every later stage is +// committed, `warpgroup_wait<1>()` retires the previous WGMMA group, and only +// then is that previous AB stage returned to the producer. The tail uses +// `warpgroup_wait<0>()` before the final release. This one-stage lag keeps TMA +// from overwriting operands still feeding a live WGMMA group. +// +// Epilogue: each N160 accumulator is staged as two N80 FP16 SMEM slices, so a +// logical N320 tile produces four N80 online-softmax folds. Both padded +// staging buffers are fully consumed before either is overwritten. Two threads +// cooperate per token row (one N40 chunk each) and merge through SMEM once the +// CTA's vocabulary split has been folded. +// ═══════════════════════════════════════════════════════════════════════════ + +// gemm_sm90.cuh carries the reusable SM90 producer/consumer mechanics and +// enables CuTe's extended MMA shapes (where the N160 WGMMA atom lives), so it +// must be the first CUTLASS/CuTe include of any TU using this header. +#include "gemm_sm90.cuh" + +#include "forward_gemm_sm90.cuh" +#include "online_softmax.cuh" + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +using namespace cute; + +// ─────────────────────────────────────────────────────────────────────────── +// Traits +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmTraitsSm90 { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + using Element = cutlass::bfloat16_t; + using ElementAccum = float; + // The N80 epilogue staging slices are FP16, mirroring the reference's + // two padded Float16 SMEM buffers. + using ElementLogit = cutlass::half_t; + + static constexpr int kTileM = Config::kTileM; + static constexpr int kTileK = Config::kTileK; + static constexpr int kAccumulatorTileN = Config::kAccumulatorTileN; + static constexpr int kLogicalTileN = Config::kLogicalTileN; + static constexpr int kEpilogueSliceN = Config::kEpilogueSliceN; + static constexpr int kSoftmaxChunkN = Config::kSoftmaxChunkN; + static constexpr int kSegmentN = Config::kSegmentN; + static constexpr int kChunksPerSegment = Config::kChunksPerSegment; + static constexpr int kStages = Config::kMainloopStages; + static constexpr int kClusterM = Config::kClusterM; + static constexpr int kNumAccumulators = Config::kNumAccumulators; + static constexpr int kNumThreads = Config::kNumThreads; + static constexpr int kWarpSize = fused_scaled_linear_cross_entropy::kWarpSize; + static constexpr int kWarpGroupSize = Config::kWarpGroupSize; + static constexpr int kNumMmaWarpGroups = Config::kNumMmaWarpGroups; + static constexpr int kConsumerThreads = Config::kConsumerThreads; + static constexpr int kLogitBuffers = Config::kLogitBuffers; + static constexpr int kLogitStride = Config::kLogitStride; + static constexpr int kRowsPerWarpGroup = Config::kRowsPerWarpGroup; + + using ClusterShape = Shape, _1, _1>; + + // hopper_helpers.make_trivial_tiled_mma(BF16, BF16, K, K, F32, (2,1,1), + // (64, ACCUMULATOR_N)): the M64 WGMMA atom, two warp groups split over M. + using TiledMmaTraits = sm90::TiledMmaMSplitSm90< + Element, + Element, + ElementAccum, + kAccumulatorTileN, + kNumMmaWarpGroups, + GMMA::Major::K, + GMMA::Major::K, + Compute>; + using TiledMma = typename TiledMmaTraits::Type; + // Pin the atom: the reference's N160 fragment must not silently decay to a + // stack of narrower WGMMA issues. + static_assert( + cute::is_same_v< + typename TiledMmaTraits::Atom, + SM90_64x160x16_F32BF16BF16_SS>, + "the forward accumulator must map onto the single N160 WGMMA atom"); + static_assert(TiledMmaTraits::kTileM == kTileM); + static_assert(TiledMmaTraits::kTileN == kAccumulatorTileN); + + // make_smem_layout_a/b for K-major BF16 with a K64 tile => 128 B rows, + // hence the SW128 swizzle atom. + using SmemTraitsX = + sm90::SmemLayoutKMajorSm90; + using SmemTraitsW = sm90::SmemLayoutKMajorSm90< + Element, kAccumulatorTileN, kTileK, kStages, Compute>; + + using SmemLayoutX1 = typename SmemTraitsX::Single; + using SmemLayoutW1 = typename SmemTraitsW::Single; + using SmemLayoutX = typename SmemTraitsX::Staged; + using SmemLayoutW = typename SmemTraitsW::Staged; + + using PipelineTraits = sm90::MainloopPipelineSm90; + using MainloopPipeline = typename PipelineTraits::Type; + using PipelineState = typename PipelineTraits::State; + + static constexpr int kTmaTransBytesX = SmemTraitsX::kStageBytes; + static constexpr int kTmaTransBytesW = SmemTraitsW::kStageBytes; + // tx_count = x_stage_bytes + NUM_ACCUMULATORS * w_stage_bytes: X and both + // N160 panels share one mbarrier per stage. + static constexpr int kTmaTransBytes = + kTmaTransBytesX + kNumAccumulators * kTmaTransBytesW; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Shared memory — same member order as the reference SharedStorage struct. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmSmemSm90 { + using Traits = ForwardGemmTraitsSm90; + using Element = typename Traits::Element; + using ElementLogit = typename Traits::ElementLogit; + + static constexpr int kSmemX = cosize_v; + static constexpr int kSmemW = cosize_v; + static constexpr int kLogitElements = + Traits::kLogitBuffers * Traits::kTileM * Traits::kLogitStride; + static constexpr int kWeightedElements = + ReturnEntropy ? Traits::kTileM : 1; + + alignas(16) typename Traits::MainloopPipeline::SharedStorage pipeline; + alignas(1024) Element smem_x[kSmemX]; + alignas(1024) Element smem_w0[kSmemW]; + alignas(1024) Element smem_w1[kSmemW]; + alignas(1024) ElementLogit smem_logits[kLogitElements]; + + // Cross-segment exchange: segment 1 publishes, segment 0 merges. + alignas(16) float segment_max[Traits::kTileM]; + alignas(16) float segment_sum[Traits::kTileM]; + alignas(16) float segment_target[Traits::kTileM]; + alignas(16) float segment_weighted[kWeightedElements]; + + CUTE_DEVICE Element* x_data() { return &smem_x[0]; } + CUTE_DEVICE Element* w0_data() { return &smem_w0[0]; } + CUTE_DEVICE Element* w1_data() { return &smem_w1[0]; } + + // s_logits[buffer, row, :] over the (LOGIT_BUFFERS, TILE_M, LOGIT_STRIDE) + // padded layout. + CUTE_DEVICE ElementLogit* logit_row(int buffer, int row) { + return &smem_logits[ + (buffer * Traits::kTileM + row) * Traits::kLogitStride]; + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// EPILOGUE_BARRIER_ID named barrier — the two consumer warp groups only. +// ─────────────────────────────────────────────────────────────────────────── + +template +CUTE_DEVICE void forward_epilogue_barrier_sm90() { + using Config = ForwardGemmConfigSm90; + sm90::named_barrier_sync_sm90< + Config::kEpilogueBarrierId, Config::kConsumerThreads, Compute>(); +} + +// ─────────────────────────────────────────────────────────────────────────── +// Work assignment — port of the kernel's cluster_work decomposition. +// ─────────────────────────────────────────────────────────────────────────── + +struct ForwardGemmWorkSm90 { + int m_pair; + int split_id; + int split_count; + int raw_pid_m; + int pid_m; + int output_work; + bool store_enabled; +}; + +template +CUTE_DEVICE ForwardGemmWorkSm90 forward_gemm_assign_work_sm90( + const ForwardGemmSplitSm90& split, + int cluster_work, + int cluster_rank) { + using Config = ForwardGemmConfigSm90; + + int m_pair; + int split_id; + int split_count; + if (split.extra_m_pairs == 0) { + m_pair = cluster_work / split.base_split_n; + split_id = cluster_work % split.base_split_n; + split_count = split.base_split_n; + } else { + // The first extra_m_pairs M pairs carry one extra vocabulary split. + int extra_work = split.extra_m_pairs * split.split_n; + if (cluster_work < extra_work) { + m_pair = cluster_work / split.split_n; + split_id = cluster_work % split.split_n; + split_count = split.split_n; + } else { + int normal_work = cluster_work - extra_work; + m_pair = split.extra_m_pairs + normal_work / split.base_split_n; + split_id = normal_work % split.base_split_n; + split_count = split.base_split_n; + } + } + + ForwardGemmWorkSm90 work; + work.m_pair = m_pair; + work.split_id = split_id; + work.split_count = split_count; + work.raw_pid_m = m_pair * Config::kClusterM + cluster_rank; + work.store_enabled = work.raw_pid_m < split.num_m_tiles; + // The tail CTA of an odd M-tile count still feeds the cluster pipeline; + // it re-reads the last valid M tile and drops its stores. + work.pid_m = work.store_enabled + ? work.raw_pid_m + : (split.num_m_tiles - 1); + work.output_work = work.raw_pid_m * split.split_n + split_id; + return work; +} + +// ─────────────────────────────────────────────────────────────────────────── +// Producer — warp 0. X is CTA-private, both W panels are cluster-multicast. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmProducerSm90 { + using Traits = ForwardGemmTraitsSm90; + static constexpr int kCompute = Compute; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + template + CUTE_DEVICE static void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + ForwardGemmSmemSm90& smem, + TmaLoadX const& tma_load_x, + TmaLoadW const& tma_load_w, + ForwardGemmParamsSm90 const& params, + ForwardGemmWorkSm90 const& work, + int num_split_tiles, + int num_k_tiles) { + // prefetch_descriptor(tma_x) / prefetch_descriptor(tma_w): warm the + // descriptor cache from the whole producer warp before gating. + cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_w.get_tma_descriptor()); + + if (cute::elect_one_sync() == 0) return; + + auto sX = make_tensor( + make_smem_ptr(smem.x_data()), typename Traits::SmemLayoutX{}); + auto sW0 = make_tensor( + make_smem_ptr(smem.w0_data()), typename Traits::SmemLayoutW{}); + auto sW1 = make_tensor( + make_smem_ptr(smem.w1_data()), typename Traits::SmemLayoutW{}); + + // int64 extents: tokens * hidden and local_vocab * hidden both exceed + // INT_MAX at production shapes and would wrap CuTe's layout algebra. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(params.tokens), + static_cast(params.hidden))); + auto mW = tma_load_w.get_tma_tensor(make_shape( + static_cast(params.local_vocab), + static_cast(params.hidden))); + + // X is not multicast; the W panels are multicast along the cluster's + // M mode, so this CTA issues its own slice of every W box. + uint32_t cluster_rank = cute::block_rank_in_cluster(); + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_w = tma_load_w.get_slice(cluster_rank); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tW0sW0 = cta_tma_w.partition_D(sW0); + auto tW1sW1 = cta_tma_w.partition_D(sW1); + + auto gX = local_tile( + mX, + make_tile(Int{}, Int{}), + make_coord(work.pid_m, _)); + auto tXgX = cta_tma_x.partition_S(gX); + + // make_layout_image_mask(cta_layout, cluster_coord, mode=0): the + // cluster is (ClusterM, 1, 1), so the W multicast covers every CTA. + uint16_t mcast_mask_w = + sm90::cluster_multicast_mask_m_sm90< + typename Traits::ClusterShape, Compute>(); + + for (int local_n = 0; local_n < num_split_tiles; ++local_n) { + int logical_n = work.split_id + local_n * work.split_count; + int fragment_base = logical_n * Traits::kNumAccumulators; + auto gW0 = local_tile( + mW, + make_tile( + Int{}, Int{}), + make_coord(fragment_base + 0, _)); + auto gW1 = local_tile( + mW, + make_tile( + Int{}, Int{}), + make_coord(fragment_base + 1, _)); + auto tW0gW0 = cta_tma_w.partition_S(gW0); + auto tW1gW1 = cta_tma_w.partition_S(gW1); + + // One mbarrier per stage covers X and both N160 panels. + for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) { + pipe.producer_acquire(state); + auto* barrier = pipe.producer_get_barrier(state); + int stage = state.index(); + copy(tma_load_x.with(*barrier, 0), + tXgX(_, _, _, k_tile), tXsX(_, _, _, stage)); + copy(tma_load_w.with(*barrier, mcast_mask_w), + tW0gW0(_, _, _, k_tile), tW0sW0(_, _, _, stage)); + copy(tma_load_w.with(*barrier, mcast_mask_w), + tW1gW1(_, _, _, k_tile), tW1sW1(_, _, _, stage)); + ++state; + } + } + + // producer_tail: keeps this CTA from retiring while a cluster peer + // still holds operands multicast out of it. + pipe.producer_tail(state); + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Consumer — warps 4..11. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmConsumerSm90 { + using Traits = ForwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Epilogue = ForwardGemmEpilogueSm90; + static constexpr int kCompute = Compute; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + // _stage_slice: scatter one N80 slice of an FP32 accumulator into a padded + // FP16 staging buffer, then publish it to both consumer warp groups. + // coord_c is thr_mma.partition_C over the identity tensor, which reproduces + // the reference's explicit row_base / col_base fragment arithmetic. + template + CUTE_DEVICE static void stage_slice( + ForwardGemmSmemSm90& smem, + Accum const& accum, + CoordC const& coord_c, + int buffer) { + CUTE_UNROLL + for (int i = 0; i < size(accum); ++i) { + int row = get<0>(coord_c(i)); + int col = get<1>(coord_c(i)); + if (col >= SliceBase && + col < SliceBase + Traits::kEpilogueSliceN) { + smem.logit_row(buffer, row)[col - SliceBase] = + static_cast(accum(i)); + } + } + cutlass::arch::fence_view_async_shared(); + asm volatile("fence.acq_rel.cta;" ::: "memory"); + forward_epilogue_barrier_sm90(); + } + + // _fold_slice: fold this thread's N40 segment of a staged N80 slice into + // the rolling statistics, then pick up the target logit if it lands here. + // `valid_cols = local_vocab - slice_vocab_base` may be negative, matching + // the reference; the chunk load keeps its mask-free fast path. + template + CUTE_DEVICE static void fold_slice( + ForwardGemmSmemSm90& smem, + OnlineSoftmaxState& softmax, + int buffer, + int row, + int segment, + int slice_vocab_base, + int valid_cols, + int target_col, + float inverse_temperature) { + constexpr int kChunkN = Traits::kSoftmaxChunkN; + constexpr int kSegmentN = Traits::kSegmentN; + int segment_base = segment * kSegmentN; + const typename Traits::ElementLogit* logits = + smem.logit_row(buffer, row); + + CUTE_UNROLL + for (int ci = 0; ci < Traits::kChunksPerSegment; ++ci) { + int c0 = segment_base + ci * kChunkN; + float values[kChunkN]; + float chunk_max = kForwardMaskLogit; + + if (c0 + kChunkN <= valid_cols) { + CUTE_UNROLL + for (int j = 0; j < kChunkN; ++j) { + float value = static_cast(logits[c0 + j]) * + inverse_temperature; + values[j] = value; + chunk_max = fmaxf(chunk_max, value); + } + } else { + CUTE_UNROLL + for (int j = 0; j < kChunkN; ++j) { + float value = static_cast(logits[c0 + j]) * + inverse_temperature; + if (c0 + j >= valid_cols) value = kForwardMaskLogit; + values[j] = value; + chunk_max = fmaxf(chunk_max, value); + } + } + + Epilogue::template fold_chunk( + softmax, values, chunk_max); + } + + // The reference re-reads the target logit from SMEM rather than + // indexing the chunk registers, which would demote them to local + // memory. target_col is already a shard-local column index. + int target_begin = slice_vocab_base + segment_base; + int target_end = target_begin + kSegmentN; + if (target_col >= target_begin && target_col < target_end) { + softmax.target_logit += + static_cast(logits[target_col - slice_vocab_base]) * + inverse_temperature; + softmax.has_target = 1; + } + } + + template + CUTE_DEVICE static void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + ForwardGemmSmemSm90& smem, + ForwardGemmParamsSm90 const& params, + ForwardGemmPartialsSm90 const& partials, + ForwardGemmWorkSm90 const& work, + int num_split_tiles, + int num_k_tiles) { + typename Traits::TiledMma tiled_mma; + int tid_in_mma = + static_cast(threadIdx.x) - Traits::kWarpGroupSize; // 0..255 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sX = make_tensor( + make_smem_ptr(smem.x_data()), typename Traits::SmemLayoutX{}); + auto sW0 = make_tensor( + make_smem_ptr(smem.w0_data()), typename Traits::SmemLayoutW{}); + auto sW1 = make_tensor( + make_smem_ptr(smem.w1_data()), typename Traits::SmemLayoutW{}); + + auto tCsX = thr_mma.partition_A(sX); + auto tCsW0 = thr_mma.partition_B(sW0); + auto tCsW1 = thr_mma.partition_B(sW1); + + using AccumShape = + Shape, Int>; + auto accum0 = partition_fragment_C(tiled_mma, AccumShape{}); + auto accum1 = partition_fragment_C(tiled_mma, AccumShape{}); + + OnlineSoftmaxState softmax; + softmax.max_value = kForwardNegInf; + auto state_release = state; + + for (int local_n = 0; local_n < num_split_tiles; ++local_n) { + int logical_n = work.split_id + local_n * work.split_count; + clear(accum0); + clear(accum1); + + // Two N160 fragments per K stage with the CuTe-DSL one-stage + // release lag: stage 0 is not released until stage 1 retires. + pipe.consumer_wait(state); + warpgroup_fence_operand(accum0); + warpgroup_fence_operand(accum1); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW0(_, _, _, state.index()), accum0); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), accum1); + warpgroup_commit_batch(); + ++state; + + for (int k_tile = 1; k_tile < num_k_tiles; ++k_tile) { + pipe.consumer_wait(state); + warpgroup_fence_operand(accum0); + warpgroup_fence_operand(accum1); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW0(_, _, _, state.index()), accum0); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), accum1); + warpgroup_commit_batch(); + + warpgroup_wait<1>(); + warpgroup_fence_operand(accum0); + warpgroup_fence_operand(accum1); + pipe.consumer_release(state_release); + ++state_release; + ++state; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(accum0); + warpgroup_fence_operand(accum1); + pipe.consumer_release(state_release); + ++state_release; + + auto coord_c = + thr_mma.partition_C(make_identity_tensor(AccumShape{})); + int tid_in_wg = tid_in_mma % Traits::kWarpGroupSize; + int mma_wg = tid_in_mma / Traits::kWarpGroupSize; + int epi_row = mma_wg * Traits::kRowsPerWarpGroup + + (tid_in_wg % Traits::kRowsPerWarpGroup); + int epi_segment = tid_in_wg / Traits::kRowsPerWarpGroup; + int global_row = work.raw_pid_m * Traits::kTileM + epi_row; + + // Keep target addressing out of the live WGMMA region. + int target_col = -1; + if (work.store_enabled && global_row < params.tokens) { + std::int64_t target_id = params.target[global_row]; + if (target_id != params.ignore_index) { + std::int64_t local = target_id - params.vocab_start; + if (local >= 0 && + local < static_cast(params.local_vocab)) { + target_col = static_cast(local); + } + } + } + + constexpr int kSliceN = Traits::kEpilogueSliceN; + constexpr int kAccumN = Traits::kAccumulatorTileN; + int logical_base = logical_n * Traits::kLogicalTileN; + int local_vocab = params.local_vocab; + + // Each N160 accumulator is folded as two N80 slices. Both buffers + // are fully consumed before either is overwritten. + stage_slice(smem, accum0, coord_c, 0); + stage_slice(smem, accum0, coord_c, 1); + fold_slice( + smem, softmax, 0, epi_row, epi_segment, + logical_base, + local_vocab - logical_base, + target_col, params.inverse_temperature); + forward_epilogue_barrier_sm90(); + + stage_slice(smem, accum1, coord_c, 0); + fold_slice( + smem, softmax, 1, epi_row, epi_segment, + logical_base + kSliceN, + local_vocab - (logical_base + kSliceN), + target_col, params.inverse_temperature); + forward_epilogue_barrier_sm90(); + + stage_slice(smem, accum1, coord_c, 1); + fold_slice( + smem, softmax, 0, epi_row, epi_segment, + logical_base + kAccumN, + local_vocab - (logical_base + kAccumN), + target_col, params.inverse_temperature); + forward_epilogue_barrier_sm90(); + + fold_slice( + smem, softmax, 1, epi_row, epi_segment, + logical_base + kAccumN + kSliceN, + local_vocab - (logical_base + kAccumN + kSliceN), + target_col, params.inverse_temperature); + forward_epilogue_barrier_sm90(); + } + + // Merge the two N40 segment halves of every token row and emit this + // CTA's split partials. + int tid_in_wg = tid_in_mma % Traits::kWarpGroupSize; + int mma_wg = tid_in_mma / Traits::kWarpGroupSize; + int epi_row = mma_wg * Traits::kRowsPerWarpGroup + + (tid_in_wg % Traits::kRowsPerWarpGroup); + int epi_segment = tid_in_wg / Traits::kRowsPerWarpGroup; + if (epi_segment == 1) { + smem.segment_max[epi_row] = softmax.max_value; + smem.segment_sum[epi_row] = softmax.exp_sum; + smem.segment_target[epi_row] = softmax.target_logit; + if constexpr (ReturnEntropy) { + smem.segment_weighted[epi_row] = softmax.exp_weighted_sum; + } + } + forward_epilogue_barrier_sm90(); + + if (epi_segment == 0 && work.store_enabled) { + OnlineSoftmaxState other; + other.max_value = smem.segment_max[epi_row]; + other.exp_sum = smem.segment_sum[epi_row]; + other.target_logit = smem.segment_target[epi_row]; + other.has_target = 0; + if constexpr (ReturnEntropy) { + other.exp_weighted_sum = smem.segment_weighted[epi_row]; + } + OnlineSoftmaxState combined = + Epilogue::template combine_scaled( + softmax, other); + Epilogue::template store_partial( + combined, + partials, + work.output_work * Traits::kTileM + epi_row); + } + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Fragment kernel +// ─────────────────────────────────────────────────────────────────────────── + +template +__global__ __launch_bounds__(ForwardGemmTraitsSm90::kNumThreads, 1) +void forward_gemm_kernel_sm90( + __grid_constant__ const TmaLoadX tma_load_x, + __grid_constant__ const TmaLoadW tma_load_w, + __grid_constant__ const ForwardGemmParamsSm90 params, + __grid_constant__ const ForwardGemmPartialsSm90 partials, + __grid_constant__ const ForwardGemmSplitSm90 split) { + using Traits = ForwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Launch = ForwardGemmLaunchSm90; + using Smem = ForwardGemmSmemSm90; + using PipelineTraits = typename Traits::PipelineTraits; + using ClusterShape = typename Traits::ClusterShape; + using Registers = sm90::WarpSpecializedRegistersSm90< + Config::kProducerRegisters, Config::kMmaRegisters, Compute>; + + extern __shared__ char raw_smem[]; + Smem& smem = *reinterpret_cast(raw_smem); + + int warp_group = cutlass::canonical_warp_group_idx(); + int warp = cutlass::canonical_warp_idx_sync(); + + // mbarrier init (stores are warp-0 only, the fence is everyone) followed by + // the cluster-wide handshake. No pipeline object exists yet, so nothing + // role-specific is live across this point. + PipelineTraits::init_barriers( + smem.pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + sm90::cluster_pipeline_init_sm90(); + + // Each role builds its own pipeline object, work assignment and tile + // counts inside a tight scope. The work/tile integers are a handful of + // scalar ops, so recomputing them per role is cheaper than keeping them + // live across the branch. + if (warp_group == 0) { + Registers::producer(); + if (warp == 0) { + auto pipe = PipelineTraits::make_producer( + smem.pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + threadIdx.x == 0, + ClusterShape{}); + typename Traits::PipelineState producer_state = + PipelineTraits::producer_start_state(); + + ForwardGemmWorkSm90 work = forward_gemm_assign_work_sm90( + split, + static_cast(blockIdx.z), + static_cast(blockIdx.x)); + int num_k_tiles = Launch::num_k_tiles(params.hidden); + int num_split_tiles = ceil_div( + split.num_logical_n_tiles - work.split_id, work.split_count); + + ForwardGemmProducerSm90::template run( + pipe, + producer_state, + smem, + tma_load_x, + tma_load_w, + params, + work, + num_split_tiles, + num_k_tiles); + } + } else { + Registers::consumer(); + auto pipe = PipelineTraits::make_consumer( + smem.pipeline, + Traits::kTmaTransBytes, + Traits::kConsumerThreads, + ClusterShape{}); + typename Traits::PipelineState consumer_state; + + ForwardGemmWorkSm90 work = forward_gemm_assign_work_sm90( + split, + static_cast(blockIdx.z), + static_cast(blockIdx.x)); + int num_k_tiles = Launch::num_k_tiles(params.hidden); + int num_split_tiles = ceil_div( + split.num_logical_n_tiles - work.split_id, work.split_count); + + ForwardGemmConsumerSm90::template run( + pipe, + consumer_state, + smem, + params, + partials, + work, + num_split_tiles, + num_k_tiles); + } + + sm90::cluster_exit_sm90(); +} + +// ─────────────────────────────────────────────────────────────────────────── +// Split reducer — the reference finalizer, stopped at *local* statistics so +// The host forward wrapper reduces across the tensor-parallel group. +// ─────────────────────────────────────────────────────────────────────────── + +template +__global__ __launch_bounds__(ForwardGemmConfigSm90::kTileM, 1) +void forward_split_reduce_kernel_sm90( + __grid_constant__ const ForwardGemmParamsSm90 params, + __grid_constant__ const ForwardGemmPartialsSm90 partials, + __grid_constant__ const ForwardGemmSplitSm90 split) { + using Config = ForwardGemmConfigSm90; + using Epilogue = ForwardGemmEpilogueSm90; + + int tid = static_cast(threadIdx.x); + int pid_m = static_cast(blockIdx.x); + int row = pid_m * Config::kTileM + tid; + if (row >= params.tokens) return; + + int split_count = split.split_count_for_pair(pid_m / Config::kClusterM); + + float row_max = kForwardNegInf; + float row_target = 0.0f; + for (int split_id = 0; split_id < split_count; ++split_id) { + int index = (pid_m * split.split_n + split_id) * Config::kTileM + tid; + row_max = fmaxf(row_max, partials.partial_max[index]); + row_target += partials.partial_target[index]; + } + + float row_sum = 0.0f; + float row_weighted = 0.0f; + for (int split_id = 0; split_id < split_count; ++split_id) { + int index = (pid_m * split.split_n + split_id) * Config::kTileM + tid; + float split_scale = forward_exp2_sm90( + (partials.partial_max[index] - row_max) * kForwardLog2E); + row_sum += partials.partial_sum[index] * split_scale; + if constexpr (ReturnEntropy) { + row_weighted += partials.partial_weighted[index] * split_scale; + } + } + + OnlineSoftmaxState state; + state.max_value = row_max; + state.exp_sum = row_sum; + state.target_logit = row_target; + if constexpr (ReturnEntropy) { + state.exp_weighted_sum = row_weighted; + } + // has_target is recomputed rather than carried: both segment threads see + // the same shard-local column, so it is a pure function of the target id. + std::int64_t target_id = params.target[row]; + std::int64_t local = target_id - params.vocab_start; + state.has_target = (target_id != params.ignore_index && local >= 0 && + local < static_cast(params.local_vocab)) ? 1 : 0; + + Epilogue::template store_row(state, params.output, row); +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90.cuh new file mode 100644 index 00000000..079152eb --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90.cuh @@ -0,0 +1,523 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy — forward GEMM contract. +// +// Structural port of +// src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_fragment_sm90.py +// (the CuTe-DSL "two-fragment" fragment kernel). Constant names below map 1:1 +// onto the reference module constants: +// +// THREADS_PER_CTA 384 TILE_M 128 TILE_K 64 STAGES 3 +// CLUSTER_M 2 ACCUMULATOR_N 160 NUM_ACCUMULATORS 2 +// LOGICAL_N 320 EPILOGUE_SLICE_N 80 CHUNK_N 40 +// LOGIT_BUFFERS 2 LOGIT_PAD 8 LOGIT_STRIDE 88 +// EPILOGUE_BARRIER_ID 4 PRODUCER_REGISTERS 24 MMA_REGISTERS 240 +// +// This header is deliberately CuTe-free: it carries the launch geometry, the +// vocabulary-split scheduler, the warp plan, the online-softmax epilogue +// contract and the host launcher declaration so torch-free consumers (and the +// existing C++ unit tests) can include the umbrella header without CUTLASS. +// +// The executable producer / consumer / mainloop is in +// forward_gemm_mainloop_sm90.cuh and is instantiated by +// fused_linear_scaled_cross_entropy_forward.cu. +// +// Tensor-parallel deltas from the single-GPU reference — everything else is a +// literal port: +// * W is the *local* vocabulary shard W_local[V_local, H]; local column c is +// global vocabulary index vocab_start + c. +// * `target` holds *global* int64 indices, mapped to the shard on the fly. +// * The reference's finalizer emits NLL/LSE/entropy directly. Here the split +// reducer stops at per-token *local* (max, sum, lse, target[, weighted]) +// statistics; fused_linear_scaled_cross_entropy_forward corrects and reduces them +// across the NVSHMEM team. +// * The reference host-pads H to a multiple of TILE_K. TMA out-of-bounds +// zero fill covers the K tail here instead, so no padded copy is made. +// H only has to keep the BF16 row stride 16 B aligned (H % 8 == 0). +// +// The [M, V] logit matrix is never materialised in HBM. +// ═══════════════════════════════════════════════════════════════════════════ + +#include + +#include +#include + +#include "config.cuh" +#include "online_softmax.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +struct ForwardLocalStatsBuffers { + float* local_max; + float* local_sum; + float* local_target; + float* local_weighted_sum; +}; + +// _fused_scaled_cross_entropy_utils_sm90.py constants. +inline constexpr float kForwardLog2E = 1.4426950408889634f; +inline constexpr float kForwardLn2 = 0.6931471805599453f; +// Rolling-max seed. Strictly greater than kForwardMaskLogit so a fold made +// entirely of masked vocabulary columns leaves the running state untouched. +inline constexpr float kForwardNegInf = -1.0e38f; +inline constexpr float kForwardMaskLogit = -3.0e38f; +inline constexpr int kHopperMaxSmemBytes = 227 * 1024; + +// exp2 with the reference's fastmath=True lowering (llvm.nvvm.ex2.approx.ftz). +__host__ __device__ inline float forward_exp2_sm90(float value) { +#if defined(__CUDA_ARCH__) + float result; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); + return result; +#else + return exp2f(value); +#endif +} + +enum class ForwardWarpRole : std::uint8_t { + kProducer, + kIdle, + kConsumer, + kInactive, +}; + +// PRODUCER_WARP_GROUP = 0 owns warp 0 (TMA). Warps 1..3 sit in the producer +// warp group but are idle in forward — they carry the dX comms in backward. +// Warps 4..11 are the two WGMMA consumer warp groups. +__host__ __device__ constexpr ForwardWarpRole forward_warp_role(int warp_id) { + return warp_id == 0 + ? ForwardWarpRole::kProducer + : (warp_id >= 1 && warp_id <= 3 + ? ForwardWarpRole::kIdle + : (warp_id >= 4 && warp_id < kNumWarps + ? ForwardWarpRole::kConsumer + : ForwardWarpRole::kInactive)); +} + +template +struct ForwardGemmConfigSm90 { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + static constexpr int kCompute = Compute; + static constexpr int kTileM = 128; + static constexpr int kTileK = 64; + static constexpr int kLogicalTileN = 320; + static constexpr int kAccumulatorTileN = 160; + static constexpr int kEpilogueSliceN = 80; + static constexpr int kSoftmaxChunkN = 40; + static constexpr int kMainloopStages = 3; + static constexpr int kClusterM = 2; + static constexpr int kNumThreads = + fused_scaled_linear_cross_entropy::kNumThreads; + + // Two N160 accumulators per consumer warp group form the logical N320. + static constexpr int kNumAccumulators = kLogicalTileN / kAccumulatorTileN; + static constexpr int kNumMmaWarpGroups = 2; + static constexpr int kWarpsPerWarpGroup = 4; + static constexpr int kWarpGroupSize = 128; + static constexpr int kConsumerThreads = + kNumMmaWarpGroups * kWarpGroupSize; + static constexpr int kMmaAtomTileM = kTileM / kNumMmaWarpGroups; + + // Two ping-pong N80 staging buffers, padded to break SMEM bank conflicts. + static constexpr int kLogitBuffers = 2; + static constexpr int kLogitPad = 8; + static constexpr int kLogitStride = kEpilogueSliceN + kLogitPad; + + // Two threads cooperate on each token row; each owns one N40 chunk of the + // N80 slice currently being folded. + static constexpr int kSegmentN = kEpilogueSliceN / 2; + static constexpr int kChunksPerSegment = kSegmentN / kSoftmaxChunkN; + static constexpr int kSegmentsPerRow = kEpilogueSliceN / kSegmentN; + static constexpr int kRowsPerWarpGroup = kTileM / kNumMmaWarpGroups; + + static constexpr int kProducerRegisters = 24; + static constexpr int kMmaRegisters = 240; + static constexpr int kUsableRegisterBudget = 64512; + + // Hardware named barrier shared by the two consumer warp groups. + static constexpr int kEpilogueBarrierId = 4; + + // _resolve_split's ceiling on the vocabulary split (config.max_split_n). + static constexpr int kDefaultMaxSplitN = 9; + + static_assert(kLogicalTileN == 2 * kAccumulatorTileN); + static_assert(kAccumulatorTileN == 2 * kEpilogueSliceN); + static_assert(kEpilogueSliceN == 2 * kSoftmaxChunkN); + static_assert(kSegmentN == kSoftmaxChunkN); + static_assert(kNumThreads == (kNumMmaWarpGroups + 1) * kWarpGroupSize); + static_assert(forward_warp_role(0) == ForwardWarpRole::kProducer); + static_assert(forward_warp_role(3) == ForwardWarpRole::kIdle); + static_assert(forward_warp_role(4) == ForwardWarpRole::kConsumer); + + __host__ __device__ static constexpr int register_total() { + return kWarpGroupSize * kProducerRegisters + + kNumMmaWarpGroups * kWarpGroupSize * kMmaRegisters; + } + static_assert(register_total() <= kUsableRegisterBudget, + "SM90 forward register budget exceeds the usable per-CTA budget"); +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Vocabulary-split scheduler — port of _resolve_split plus the kernel-side +// cluster_work decomposition. +// ─────────────────────────────────────────────────────────────────────────── + +// ScaledCEForwardFragmentConfig's tuning knobs. Zero means "auto". +template +struct ForwardGemmTuningSm90 { + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + int split_n = 0; + int base_split_n = 0; + int extra_m_pairs = 0; + int target_cluster_pairs = 0; + int max_split_n = Config::kDefaultMaxSplitN; +}; + +template +struct ForwardGemmSplitSm90 { + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + int split_n = 1; + int base_split_n = 1; + int extra_m_pairs = 0; + int num_m_tiles = 0; + int num_m_pairs = 0; + int num_logical_n_tiles = 0; + int num_cluster_pairs = 0; + + // Number of vocabulary splits that actually cover `m_pair`. M pairs below + // extra_m_pairs get the extra split; the rest get base_split_n. + __host__ __device__ int split_count_for_pair(int m_pair) const { + if (extra_m_pairs == 0) return base_split_n; + return m_pair < extra_m_pairs ? split_n : base_split_n; + } +}; + +// FP32 (max, sum, target[, weighted]) split partials, laid out +// [num_m_tiles * split_n, TILE_M] exactly like the reference's +// partial_* tensors. +template +struct ForwardGemmPartialsSm90 { + static constexpr int kCompute = Compute; + + float* partial_max = nullptr; + float* partial_sum = nullptr; + float* partial_target = nullptr; + float* partial_weighted = nullptr; +}; + +// Problem + output description. The vocabulary shard is contiguous: local +// column c is global vocabulary index `vocab_start + c`, and targets arrive as +// *global* int64 indices. +template +struct ForwardGemmParamsSm90 { + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + const void* x = nullptr; // BF16 [tokens, hidden], row major + const void* weight = nullptr; // BF16 [local_vocab, hidden], row major + const std::int64_t* target = nullptr; // int64 [tokens], global indices + ForwardLocalStatsBuffers output = {}; // FP32 per-token local statistics + + // Split partials scratch; size it with + // ForwardGemmLaunchSm90::workspace_bytes(...). + void* workspace = nullptr; + std::size_t workspace_bytes = 0; + + int tokens = 0; + int hidden = 0; + int local_vocab = 0; + std::int64_t vocab_start = 0; + std::int64_t ignore_index = -100; + float inverse_temperature = 1.0f; + + ForwardGemmTuningSm90 tuning = {}; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Online-softmax epilogue contract +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmEpilogueSm90 { + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + // Scalar reference fold: one online-softmax update per logit. Kept as the + // readable contract; the device path uses the chunked form below. + template + __host__ __device__ static void fold_fragment( + OnlineSoftmaxState& state, + const float (&scaled_logits)[FragmentSize], + int fragment_vocab_start, + int local_vocab_size, + int local_target) { + #pragma unroll + for (int i = 0; i < FragmentSize; ++i) { + int local_vocab_index = fragment_vocab_start + i; + state.template fold( + scaled_logits[i], + local_vocab_index < local_vocab_size, + local_vocab_index == local_target); + } + } + + // _fold_slice's rolling update, verbatim: one rescale of the running + // (sum[, weighted]) per N40 chunk. `values` is the chunk already masked + // and scaled by the caller and `chunk_max` its maximum, matching the + // reference's two-branch load. + template + __host__ __device__ static void fold_chunk( + OnlineSoftmaxState& state, + const float (&values)[ChunkN], + float chunk_max) { + float next_max = fmaxf(state.max_value, chunk_max); + float chunk_sum = 0.0f; + + if constexpr (ReturnEntropy) { + float chunk_weighted = 0.0f; + #pragma unroll + for (int i = 0; i < ChunkN; ++i) { + float weight = + forward_exp2_sm90((values[i] - next_max) * kForwardLog2E); + chunk_sum += weight; + chunk_weighted += weight * values[i]; + } + float previous_scale = forward_exp2_sm90( + (state.max_value - next_max) * kForwardLog2E); + state.exp_sum = state.exp_sum * previous_scale + chunk_sum; + state.exp_weighted_sum = + state.exp_weighted_sum * previous_scale + chunk_weighted; + } else { + #pragma unroll + for (int i = 0; i < ChunkN; ++i) { + chunk_sum += + forward_exp2_sm90((values[i] - next_max) * kForwardLog2E); + } + state.exp_sum = state.exp_sum * forward_exp2_sm90( + (state.max_value - next_max) * kForwardLog2E) + chunk_sum; + } + + state.max_value = next_max; + } + + // Branch-free merge of two rolling states, matching the reference's + // cross-segment combine and its split finalizer. Deliberately not + // merge_online_softmax(): that shortcuts on exp_sum == 0 and would put + // data-dependent branches in the epilogue. + template + __host__ __device__ static OnlineSoftmaxState combine_scaled( + const OnlineSoftmaxState& lhs, const OnlineSoftmaxState& rhs) { + OnlineSoftmaxState result; + result.max_value = fmaxf(lhs.max_value, rhs.max_value); + float lhs_scale = forward_exp2_sm90( + (lhs.max_value - result.max_value) * kForwardLog2E); + float rhs_scale = forward_exp2_sm90( + (rhs.max_value - result.max_value) * kForwardLog2E); + result.exp_sum = lhs.exp_sum * lhs_scale + rhs.exp_sum * rhs_scale; + if constexpr (ReturnEntropy) { + result.exp_weighted_sum = + lhs.exp_weighted_sum * lhs_scale + + rhs.exp_weighted_sum * rhs_scale; + } + // The reference accumulates the target logit additively: at most one + // contributor is non-zero because the shard is partitioned. + result.target_logit = lhs.target_logit + rhs.target_logit; + result.has_target = lhs.has_target || rhs.has_target; + return result; + } + + template + __host__ __device__ static void store_partial( + const OnlineSoftmaxState& state, + const ForwardGemmPartialsSm90& partials, + int index) { + partials.partial_max[index] = state.max_value; + partials.partial_sum[index] = state.exp_sum; + partials.partial_target[index] = state.target_logit; + if constexpr (ReturnEntropy) { + partials.partial_weighted[index] = state.exp_weighted_sum; + } + } + + template + __host__ __device__ static void store_row( + const OnlineSoftmaxState& state, + const ForwardLocalStatsBuffers& output, + int row) { + output.local_max[row] = state.max_value; + output.local_sum[row] = state.exp_sum; + output.local_target[row] = state.has_target ? state.target_logit : 0.0f; + if constexpr (ReturnEntropy) { + output.local_weighted_sum[row] = state.exp_weighted_sum; + } + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Host-side launch geometry. CuTe-free so the launcher declaration, the split +// resolution and the shared-memory budget stay visible without CUTLASS. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ForwardGemmLaunchSm90 { + using Config = ForwardGemmConfigSm90; + static constexpr int kCompute = Compute; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + __host__ __device__ static constexpr int num_m_tiles(int tokens) { + return ceil_div(tokens, Config::kTileM); + } + + __host__ __device__ static constexpr int num_m_pairs(int tokens) { + return ceil_div(num_m_tiles(tokens), Config::kClusterM); + } + + __host__ __device__ static constexpr int num_k_tiles(int hidden) { + return ceil_div(hidden, Config::kTileK); + } + + __host__ __device__ static constexpr int num_logical_n_tiles( + int local_vocab) { + return ceil_div(local_vocab, Config::kLogicalTileN); + } + + // _resolve_split. `max_active_cluster_pairs` is the device's occupancy for + // a CLUSTER_M cluster (cutlass HardwareInfo.get_max_active_clusters). + __host__ static ForwardGemmSplitSm90 resolve_split( + const ForwardGemmTuningSm90& tuning, + int tokens, + int local_vocab, + int max_active_cluster_pairs) { + ForwardGemmSplitSm90 split; + split.num_m_tiles = num_m_tiles(tokens); + split.num_m_pairs = num_m_pairs(tokens); + split.num_logical_n_tiles = num_logical_n_tiles(local_vocab); + + if (tuning.base_split_n != 0) { + split.split_n = tuning.split_n; + split.base_split_n = tuning.base_split_n; + split.extra_m_pairs = tuning.extra_m_pairs; + } else if (tuning.split_n != 0) { + split.split_n = tuning.split_n; + split.base_split_n = tuning.split_n; + split.extra_m_pairs = 0; + } else { + int max_split_n = tuning.max_split_n > 0 + ? tuning.max_split_n + : Config::kDefaultMaxSplitN; + int target_cluster_pairs = tuning.target_cluster_pairs > 0 + ? tuning.target_cluster_pairs + : (max_active_cluster_pairs > 0 ? max_active_cluster_pairs : 1); + int per_pair_cap = max_split_n < split.num_logical_n_tiles + ? max_split_n + : split.num_logical_n_tiles; + int max_cluster_pairs = split.num_m_pairs * per_pair_cap; + int wanted = target_cluster_pairs < max_cluster_pairs + ? target_cluster_pairs + : max_cluster_pairs; + int cluster_pairs = + split.num_m_pairs > wanted ? split.num_m_pairs : wanted; + split.base_split_n = cluster_pairs / split.num_m_pairs; + split.extra_m_pairs = cluster_pairs % split.num_m_pairs; + split.split_n = + split.base_split_n + (split.extra_m_pairs != 0 ? 1 : 0); + } + + split.num_cluster_pairs = + split.num_m_pairs * split.base_split_n + split.extra_m_pairs; + return split; + } + + // Upper bound on the split-partial scratch, independent of the runtime + // split resolution so callers can size it ahead of the launch. + __host__ __device__ static std::size_t workspace_bytes( + int tokens, + int local_vocab, + bool return_entropy, + int max_split_n = Config::kDefaultMaxSplitN) { + int logical_n = num_logical_n_tiles(local_vocab); + if (logical_n < 1) logical_n = 1; + int cap = max_split_n < logical_n ? max_split_n : logical_n; + if (cap < 1) cap = 1; + std::size_t rows = static_cast(num_m_tiles(tokens)) * + static_cast(cap) * + static_cast(Config::kTileM); + return rows * sizeof(float) * (return_entropy ? 4u : 3u); + } + + // Mirrors ForwardGemmSmemSm90 — duplicated so + // callers can validate the dynamic allocation without CUTLASS. + __host__ __device__ static constexpr int smem_bytes(bool return_entropy) { + return + // X + two N160 weight panels, BF16, kMainloopStages deep. + (Config::kTileM + Config::kLogicalTileN) * Config::kTileK * 2 * + Config::kMainloopStages + + // Two padded FP16 N80 staging buffers. + Config::kLogitBuffers * Config::kTileM * Config::kLogitStride * 2 + + // Cross-segment (max, sum, target) exchange. + 3 * Config::kTileM * 4 + + // Optional weighted-entropy moment exchange. + (return_entropy ? Config::kTileM : 1) * 4 + + // TMA mbarrier pairs. + 16 * Config::kMainloopStages + + // Struct alignment padding (four 1024 B aligned regions). + 4 * 1024; + } + + static_assert(smem_bytes(true) <= kHopperMaxSmemBytes, + "SM90 forward shared-memory footprint exceeds the Hopper limit"); +}; + +// End-to-end tensor-parallel forward: the exact CuTe-DSL local-statistic GEMM +// followed by NCCL MAX, correction/packing, NCCL SUM, and finalization. This +// is the sole production forward host launcher; it directly prepares and +// launches forward_gemm_kernel_sm90. +// +// `gemm` describes this rank's contiguous vocabulary shard. Its `output` and +// `workspace` fields are launcher-owned and overwritten: both the local +// statistics and the vocabulary-split partials are internal scratch that never +// escapes the operation, so they come from liger::global_buffer_pool() rather +// than from the caller. Call configure_forward_tp_workspace() first (see +// forward_reduce.cuh). +template +struct ForwardTpParamsSm90 { + static constexpr int kCompute = Compute; + + ForwardGemmParamsSm90 gemm = {}; + + float* nll = nullptr; // [tokens] + float* lse = nullptr; // [tokens] + float* entropy = nullptr; // [tokens], required when ReturnEntropy=true + + // ncclComm_t for the tensor-parallel communicator, as an int64 handle. + std::int64_t nccl_comm_handle = 0; +}; + +template +void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90& params, + cudaStream_t stream); + +extern template void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90<90>&, + cudaStream_t); +extern template void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90<90>&, + cudaStream_t); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90_selftest.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90_selftest.cu new file mode 100644 index 00000000..a3fe12a9 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90_selftest.cu @@ -0,0 +1,414 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// Internal SM90 forward GEMM self test. +// +// Compiled out unless LIGER_FSLCE_FORWARD_SM90_SELFTEST is defined, so the +// core library's recursive .cu glob picks this file up as an empty translation +// unit and no shared CMake/test file has to change. Build it standalone with: +// +// nvcc -std=c++17 -arch=sm_90a -O3 --expt-relaxed-constexpr \ +// -DLIGER_FSLCE_FORWARD_SM90_SELFTEST \ +// forward_gemm_sm90_selftest.cu \ +// -I -I -I$CUTLASS_HOME/include -o selftest +// +// Checks the fused kernel's local (max, sum, lse, target, weighted) statistics +// against a CPU reference that reproduces the same staging precision: the +// reference rounds every logit through FP16 exactly like the N80 SMEM slices, +// then applies the temperature scaling. +// ═══════════════════════════════════════════════════════════════════════════ + +#if defined(LIGER_FSLCE_FORWARD_SM90_SELFTEST) + +#include "forward_gemm_mainloop_sm90.cuh" + +#include "forward_gemm_sm90.cuh" + +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +namespace fslce = liger::fused_scaled_linear_cross_entropy; + +using Element = cutlass::bfloat16_t; +using ElementLogit = cutlass::half_t; + +int g_failures = 0; + +#define CUDA_OK(expr) \ + do { \ + cudaError_t status = (expr); \ + if (status != cudaSuccess) { \ + std::printf("CUDA error %s at %s:%d: %s\n", \ + #expr, __FILE__, __LINE__, cudaGetErrorString(status)); \ + std::exit(1); \ + } \ + } while (0) + +float next_random(std::uint64_t& state) { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + std::uint32_t bits = static_cast(state >> 33); + return static_cast(bits % 2048) / 1024.0f - 1.0f; +} + +struct ReferenceRow { + float max_value; + float exp_sum; + float lse; + float target_logit; + float weighted_sum; +}; + +// CPU mirror of the device epilogue: FP32 dot product, FP16 staging round, +// temperature scaling, then a plain (non-online) softmax reduction. +std::vector reference( + const std::vector& x, + const std::vector& weight, + const std::vector& target, + int tokens, + int hidden, + int local_vocab, + std::int64_t vocab_start, + std::int64_t ignore_index, + float inverse_temperature) { + std::vector rows(tokens); + std::vector scaled(local_vocab); + for (int m = 0; m < tokens; ++m) { + float row_max = -1.0e38f; + for (int v = 0; v < local_vocab; ++v) { + float dot = 0.0f; + for (int k = 0; k < hidden; ++k) { + dot += static_cast(x[m * hidden + k]) * + static_cast(weight[v * hidden + k]); + } + float staged = static_cast(static_cast(dot)); + scaled[v] = staged * inverse_temperature; + row_max = fmaxf(row_max, scaled[v]); + } + + double sum = 0.0; + double weighted = 0.0; + for (int v = 0; v < local_vocab; ++v) { + double e = std::exp( + static_cast(scaled[v]) - static_cast(row_max)); + sum += e; + weighted += e * static_cast(scaled[v]); + } + + float target_logit = 0.0f; + std::int64_t target_id = target[m]; + if (target_id != ignore_index) { + std::int64_t local = target_id - vocab_start; + if (local >= 0 && local < local_vocab) { + target_logit = scaled[static_cast(local)]; + } + } + + rows[m].max_value = row_max; + rows[m].exp_sum = static_cast(sum); + rows[m].lse = row_max + static_cast(std::log(sum)); + rows[m].target_logit = target_logit; + rows[m].weighted_sum = static_cast(weighted); + } + return rows; +} + +void expect_near( + const char* label, + int row, + float actual, + float expected, + float absolute, + float relative) { + float tolerance = absolute + relative * fabsf(expected); + if (!(fabsf(actual - expected) <= tolerance)) { + std::printf( + " FAIL %-14s row %4d: got %14.6f want %14.6f (tol %.6f)\n", + label, row, actual, expected, tolerance); + ++g_failures; + } +} + +// Test-only direct kernel launch. Production has no GEMM-only host wrapper: +// fused_linear_scaled_cross_entropy_forward owns this setup. +template +void launch_forward_gemm_selftest( + const fslce::ForwardGemmParamsSm90<90>& params, + cudaStream_t stream) { + using Traits = fslce::ForwardGemmTraitsSm90<90>; + using Config = typename Traits::Config; + using Launch = fslce::ForwardGemmLaunchSm90<90>; + using Smem = fslce::ForwardGemmSmemSm90<90, ReturnEntropy>; + using Cluster = liger::sm90::ClusterLaunchSm90<90>; + + auto tma_load_x = liger::sm90::TmaSm90<90>::template make_load< + typename Traits::SmemLayoutX1, + Traits::kTileM, + Traits::kTileK>( + static_cast(params.x), + params.tokens, + params.hidden); + auto tma_load_w = + liger::sm90::TmaSm90<90>::template make_load_multicast< + typename Traits::SmemLayoutW1, + Traits::kAccumulatorTileN, + Traits::kTileK, + Traits::kClusterM>( + static_cast(params.weight), + params.local_vocab, + params.hidden); + auto* kernel_fn = &fslce::forward_gemm_kernel_sm90< + ReturnEntropy, 90, decltype(tma_load_x), decltype(tma_load_w)>; + + constexpr int kSmemBytes = static_cast(sizeof(Smem)); + CUDA_OK(Cluster::prepare(kernel_fn, kSmemBytes)); + int max_active_cluster_pairs = Cluster::max_active_clusters( + kernel_fn, Config::kNumThreads, kSmemBytes, Config::kClusterM); + fslce::ForwardGemmSplitSm90<90> split = Launch::resolve_split( + params.tuning, + params.tokens, + params.local_vocab, + max_active_cluster_pairs); + + std::size_t partial_rows = + static_cast(split.num_m_tiles) * + static_cast(split.split_n) * + static_cast(Config::kTileM); + float* scratch = static_cast(params.workspace); + fslce::ForwardGemmPartialsSm90<90> partials; + partials.partial_max = scratch + 0 * partial_rows; + partials.partial_sum = scratch + 1 * partial_rows; + partials.partial_target = scratch + 2 * partial_rows; + partials.partial_weighted = + ReturnEntropy ? scratch + 3 * partial_rows : nullptr; + + CUDA_OK(Cluster::launch( + kernel_fn, + dim3( + static_cast(Config::kClusterM), + 1u, + static_cast(split.num_cluster_pairs)), + Config::kNumThreads, + kSmemBytes, + Config::kClusterM, + stream, + tma_load_x, + tma_load_w, + params, + partials, + split)); + fslce::forward_split_reduce_kernel_sm90 + <<>>( + params, partials, split); +} + +template +void run_case( + const std::string& name, + int tokens, + int hidden, + int local_vocab, + std::int64_t vocab_start, + float temperature, + std::uint64_t seed, + fslce::ForwardGemmTuningSm90<90> tuning = {}) { + std::uint64_t state = seed; + std::vector host_x(static_cast(tokens) * hidden); + std::vector host_w(static_cast(local_vocab) * hidden); + std::vector host_target(tokens); + + const float scale = 1.0f / std::sqrt(static_cast(hidden)); + for (auto& value : host_x) { + value = static_cast(next_random(state) * scale * 4.0f); + } + for (auto& value : host_w) { + value = static_cast(next_random(state) * scale * 4.0f); + } + const std::int64_t global_vocab = vocab_start + local_vocab + 37; + for (int m = 0; m < tokens; ++m) { + if (m % 11 == 3) { + host_target[m] = -100; // ignore_index + } else if (m % 5 == 0) { + // Force a target that lands outside this vocabulary shard. + host_target[m] = (vocab_start + local_vocab) % global_vocab; + } else { + std::uint64_t bits = state = state * 2862933555777941757ULL + 3037000493ULL; + host_target[m] = vocab_start + + static_cast(bits % static_cast(local_vocab)); + } + } + + Element* device_x = nullptr; + Element* device_w = nullptr; + std::int64_t* device_target = nullptr; + CUDA_OK(cudaMalloc(&device_x, host_x.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_w, host_w.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_target, host_target.size() * sizeof(std::int64_t))); + CUDA_OK(cudaMemcpy(device_x, host_x.data(), + host_x.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_w, host_w.data(), + host_w.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_target, host_target.data(), + host_target.size() * sizeof(std::int64_t), cudaMemcpyHostToDevice)); + + float* device_stats = nullptr; + CUDA_OK(cudaMalloc(&device_stats, 4 * tokens * sizeof(float))); + CUDA_OK(cudaMemset(device_stats, 0xff, 4 * tokens * sizeof(float))); + + using Launch = fslce::ForwardGemmLaunchSm90<90>; + std::size_t workspace_bytes = Launch::workspace_bytes( + tokens, local_vocab, ReturnEntropy, tuning.max_split_n); + void* device_workspace = nullptr; + CUDA_OK(cudaMalloc(&device_workspace, workspace_bytes)); + + fslce::ForwardGemmParamsSm90<90> params; + params.x = device_x; + params.weight = device_w; + params.target = device_target; + params.output.local_max = device_stats + 0 * tokens; + params.output.local_sum = device_stats + 1 * tokens; + params.output.local_target = device_stats + 2 * tokens; + params.output.local_weighted_sum = device_stats + 3 * tokens; + params.workspace = device_workspace; + params.workspace_bytes = workspace_bytes; + params.tokens = tokens; + params.hidden = hidden; + params.local_vocab = local_vocab; + params.vocab_start = vocab_start; + params.ignore_index = -100; + params.inverse_temperature = 1.0f / temperature; + params.tuning = tuning; + + launch_forward_gemm_selftest(params, nullptr); + CUDA_OK(cudaGetLastError()); + CUDA_OK(cudaDeviceSynchronize()); + + std::vector host_stats(4 * tokens); + CUDA_OK(cudaMemcpy(host_stats.data(), device_stats, + host_stats.size() * sizeof(float), cudaMemcpyDeviceToHost)); + + std::vector expected = reference( + host_x, host_w, host_target, tokens, hidden, local_vocab, + vocab_start, params.ignore_index, params.inverse_temperature); + + int before = g_failures; + for (int m = 0; m < tokens; ++m) { + expect_near("local_max", m, host_stats[0 * tokens + m], + expected[m].max_value, 2.0e-3f, 1.0e-3f); + expect_near("local_sum", m, host_stats[1 * tokens + m], + expected[m].exp_sum, 1.0e-3f, 5.0e-3f); + expect_near("local_target", m, host_stats[2 * tokens + m], + expected[m].target_logit, 2.0e-3f, 1.0e-3f); + if constexpr (ReturnEntropy) { + expect_near("local_weighted", m, host_stats[3 * tokens + m], + expected[m].weighted_sum, 5.0e-3f, 5.0e-3f); + } + if (g_failures > before + 20) break; + } + + std::printf("%-46s %s (M=%d H=%d Vloc=%d start=%lld T=%.2f entropy=%d " + "tuning split=%d base=%d extra=%d)\n", + name.c_str(), + g_failures == before ? "PASS" : "FAIL", + tokens, hidden, local_vocab, + static_cast(vocab_start), + temperature, + static_cast(ReturnEntropy), + tuning.split_n, tuning.base_split_n, tuning.extra_m_pairs); + + CUDA_OK(cudaFree(device_x)); + CUDA_OK(cudaFree(device_w)); + CUDA_OK(cudaFree(device_target)); + CUDA_OK(cudaFree(device_stats)); + CUDA_OK(cudaFree(device_workspace)); +} + +void report_occupancy() { + using Traits = fslce::ForwardGemmTraitsSm90<90>; + std::printf( + "geometry: threads=%d cluster_m=%d stages=%d tileM=%d tileK=%d " + "logicalN=%d accumN=%d sliceN=%d chunkN=%d tma_bytes=%d\n", + Traits::kNumThreads, Traits::kClusterM, Traits::kStages, + Traits::kTileM, Traits::kTileK, Traits::kLogicalTileN, + Traits::kAccumulatorTileN, Traits::kEpilogueSliceN, + Traits::kSoftmaxChunkN, Traits::kTmaTransBytes); + std::printf("smem: no-entropy=%zu B entropy=%zu B (Hopper limit %d B)\n", + sizeof(fslce::ForwardGemmSmemSm90<90, false>), + sizeof(fslce::ForwardGemmSmemSm90<90, true>), + fslce::kHopperMaxSmemBytes); +} + +} // namespace + +int main() { + int device = 0; + cudaDeviceProp properties = {}; + CUDA_OK(cudaGetDeviceProperties(&properties, device)); + if (properties.major != 9) { + std::printf("skipping: requires a Hopper (sm90) GPU, found sm_%d%d\n", + properties.major, properties.minor); + return 0; + } + report_occupancy(); + + try { + // Aligned: M and V are exact multiples of the tile geometry. + run_case("aligned / no entropy", 256, 128, 640, 0, 1.0f, 1); + run_case("aligned / entropy", 256, 128, 640, 0, 1.0f, 2); + // Ragged M (odd M-tile count exercises the redundant cluster peer), + // ragged H (TMA zero fill in K) and ragged V (epilogue masking). + // H stays a multiple of 8 so the BF16 row stride keeps TMA's 16 B + // global-alignment requirement. + run_case("ragged M/H/V / no entropy", 200, 104, 517, 0, 1.0f, 3); + run_case("ragged M/H/V / entropy", 200, 104, 517, 0, 1.0f, 4); + // K tail: 72 = one full K64 tile plus 8 valid columns. + run_case("ragged K tail", 128, 72, 333, 0, 1.0f, 8); + // Sharded vocabulary with a non-zero start plus temperature scaling. + run_case("shard offset + temperature", 384, 192, 901, 4096, 0.7f, 5); + run_case("single M tile", 96, 64, 321, 12288, 1.3f, 6); + // Local vocabulary smaller than one logical N320 tile. + run_case("tiny local vocabulary", 128, 64, 40, 100000, 1.0f, 7); + + // Vocabulary-split scheduler. M=1024 -> 8 M tiles -> 4 M pairs; + // V=2560 -> 8 logical N320 tiles. + fslce::ForwardGemmTuningSm90<90> uniform; + uniform.split_n = 3; + uniform.base_split_n = 3; + uniform.extra_m_pairs = 0; + run_case("split: uniform base_split_n=3", + 1024, 128, 2560, 0, 1.0f, 9, uniform); + + // Uneven split: the first 3 of 4 M pairs carry an extra vocabulary + // split, so cluster_work crosses both scheduler branches. + fslce::ForwardGemmTuningSm90<90> uneven; + uneven.split_n = 3; + uneven.base_split_n = 2; + uneven.extra_m_pairs = 3; + run_case("split: uneven base=2 extra=3", + 1024, 128, 2560, 0, 1.0f, 10, uneven); + run_case("split: uneven / no entropy", + 1024, 128, 2560, 7168, 0.85f, 11, uneven); + + // Max split with a ragged M tail on top of an uneven schedule. + fslce::ForwardGemmTuningSm90<90> ragged_split; + ragged_split.split_n = 2; + ragged_split.base_split_n = 1; + ragged_split.extra_m_pairs = 1; + run_case("split: uneven + ragged M/V", + 648, 136, 1613, 512, 1.0f, 12, ragged_split); + } catch (const std::exception& error) { + std::printf("exception: %s\n", error.what()); + return 1; + } + + std::printf("%s: %d comparison failures\n", + g_failures == 0 ? "OK" : "FAILED", g_failures); + return g_failures == 0 ? 0 : 1; +} + +#endif // LIGER_FSLCE_FORWARD_SM90_SELFTEST diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cu new file mode 100644 index 00000000..9e6b0172 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cu @@ -0,0 +1,132 @@ +#include "forward_reduce.cuh" + +#include + +#include "buffer_pool.cuh" +#include "liger_cute/check.h" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { +namespace { + +int g_capacity_tokens = 0; +std::size_t g_capacity_split_partials_bytes = 0; + +std::size_t token_bytes_at(int tokens) { + return static_cast(tokens) * sizeof(float); +} + +std::size_t packed_bytes_at(int tokens) { + return static_cast(kForwardReducedFields) * + token_bytes_at(tokens); +} + +std::size_t split_partials_bytes_at(int tokens, int local_vocab) { + return ForwardGemmLaunchSm90<90>::workspace_bytes( + tokens, local_vocab, /*return_entropy=*/true); +} + +void ensure_capacity(int tokens) { + LIGER_CHECK(tokens > 0, "tokens must be positive"); + LIGER_CHECK( + g_capacity_tokens != 0, + "fused_scaled_linear_cross_entropy forward: call " + "configure_forward_tp_workspace(max_tokens, max_local_vocab) before " + "the first launch"); + LIGER_CHECK( + tokens <= g_capacity_tokens, + "fused_scaled_linear_cross_entropy forward workspace was configured " + "for ", + g_capacity_tokens, + " tokens but ", + tokens, + " were requested"); +} + +} // namespace + +void configure_forward_tp_workspace(int max_tokens, int max_local_vocab) { + LIGER_CHECK(max_tokens > 0, "max_tokens must be positive"); + LIGER_CHECK(max_local_vocab > 0, "max_local_vocab must be positive"); + + std::size_t split_bytes = + split_partials_bytes_at(max_tokens, max_local_vocab); + if (max_tokens <= g_capacity_tokens && + split_bytes <= g_capacity_split_partials_bytes) { + return; + } + LIGER_CHECK( + g_capacity_tokens == 0, + "fused_scaled_linear_cross_entropy forward workspace is immutable " + "after its first configuration"); + + using Names = ForwardBufferNames; + std::size_t token_bytes = token_bytes_at(max_tokens); + std::size_t packed_bytes = packed_bytes_at(max_tokens); + auto& pool = global_buffer_pool(); + + pool.get_device(Names::kLocalMax, token_bytes); + pool.get_device(Names::kGlobalMax, token_bytes); + pool.get_device(Names::kPacked, packed_bytes); + pool.get_device(Names::kReduced, packed_bytes); + pool.get_device(Names::kGemmSplitPartials, split_bytes); + pool.get_device(Names::kLocalSum, token_bytes); + pool.get_device(Names::kLocalTarget, token_bytes); + pool.get_device(Names::kLocalWeighted, token_bytes); + + g_capacity_tokens = max_tokens; + g_capacity_split_partials_bytes = split_bytes; +} + +std::size_t forward_tp_workspace_device_bytes( + int max_tokens, int max_local_vocab) { + LIGER_CHECK(max_tokens > 0, "max_tokens must be positive"); + LIGER_CHECK(max_local_vocab > 0, "max_local_vocab must be positive"); + return 5 * token_bytes_at(max_tokens) + + 2 * packed_bytes_at(max_tokens) + + split_partials_bytes_at(max_tokens, max_local_vocab); +} + +void reset_forward_tp_workspace_configuration() { + g_capacity_tokens = 0; + g_capacity_split_partials_bytes = 0; +} + +template +ForwardTpWorkspace reserve_forward_tp_workspace(int tokens) { + ensure_capacity(tokens); + + using Names = ForwardBufferNames; + std::size_t token_bytes = token_bytes_at(g_capacity_tokens); + std::size_t packed_bytes = packed_bytes_at(g_capacity_tokens); + auto& pool = global_buffer_pool(); + + ForwardTpWorkspace workspace = {}; + workspace.local.local_max = static_cast( + pool.get_device(Names::kLocalMax, token_bytes)); + workspace.global_max = static_cast( + pool.get_device(Names::kGlobalMax, token_bytes)); + workspace.packed = static_cast( + pool.get_device(Names::kPacked, packed_bytes)); + workspace.reduced = static_cast( + pool.get_device(Names::kReduced, packed_bytes)); + workspace.gemm_split_partials = pool.get_device( + Names::kGemmSplitPartials, g_capacity_split_partials_bytes); + workspace.gemm_split_partials_bytes = + g_capacity_split_partials_bytes; + workspace.local.local_sum = static_cast( + pool.get_device(Names::kLocalSum, token_bytes)); + workspace.local.local_target = static_cast( + pool.get_device(Names::kLocalTarget, token_bytes)); + workspace.local.local_weighted_sum = static_cast( + pool.get_device(Names::kLocalWeighted, token_bytes)); + workspace.tokens = tokens; + workspace.fields = forward_reduced_fields(ReturnEntropy); + return workspace; +} + +template ForwardTpWorkspace reserve_forward_tp_workspace(int); +template ForwardTpWorkspace reserve_forward_tp_workspace(int); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cuh new file mode 100644 index 00000000..c57bf303 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cuh @@ -0,0 +1,80 @@ +#pragma once + +// Launcher-owned workspace for the tensor-parallel forward. The GEMM and +// split reducer produce rank-local FP32 statistics here; the host launcher then +// performs NCCL MAX/SUM all-reduces with small CUDA kernels between/after them. +// No forward buffer is remotely addressed by device code, so every allocation +// is ordinary device memory from the shared BufferPool. + +#include + +#include + +#include "forward_gemm_sm90.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +inline constexpr int kForwardReducedSumField = 0; +inline constexpr int kForwardReducedTargetField = 1; +inline constexpr int kForwardReducedWeightedField = 2; +inline constexpr int kForwardReducedFields = 3; + +struct ForwardBufferNames { + static constexpr const char* kLocalMax = + "fused_scaled_linear_cross_entropy_tp_forward_local_max"; + static constexpr const char* kGlobalMax = + "fused_scaled_linear_cross_entropy_tp_forward_global_max"; + static constexpr const char* kPacked = + "fused_scaled_linear_cross_entropy_tp_forward_packed"; + static constexpr const char* kReduced = + "fused_scaled_linear_cross_entropy_tp_forward_reduced"; + static constexpr const char* kGemmSplitPartials = + "fused_scaled_linear_cross_entropy_tp_forward_split_partials"; + static constexpr const char* kLocalSum = + "fused_scaled_linear_cross_entropy_tp_forward_local_sum"; + static constexpr const char* kLocalTarget = + "fused_scaled_linear_cross_entropy_tp_forward_local_target"; + static constexpr const char* kLocalWeighted = + "fused_scaled_linear_cross_entropy_tp_forward_local_weighted"; +}; + +__host__ __device__ constexpr int forward_reduced_fields( + bool return_entropy) { + return return_entropy + ? kForwardReducedFields + : kForwardReducedFields - 1; +} + +struct ForwardTpWorkspace { + ForwardLocalStatsBuffers local; + float* global_max; + float* packed; + float* reduced; + void* gemm_split_partials; + std::size_t gemm_split_partials_bytes; + int tokens; + int fields; +}; + +// Fixes reusable workspace capacities. Capacities are immutable after the +// first call; callers must configure their maximum shape upfront. +void configure_forward_tp_workspace(int max_tokens, int max_local_vocab); + +std::size_t forward_tp_workspace_device_bytes( + int max_tokens, + int max_local_vocab); + +// Invalidates cached capacities after the shared BufferPool is cleared. +void reset_forward_tp_workspace_configuration(); + +template +ForwardTpWorkspace reserve_forward_tp_workspace(int tokens); + +extern template ForwardTpWorkspace + reserve_forward_tp_workspace(int); +extern template ForwardTpWorkspace + reserve_forward_tp_workspace(int); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_tp_sm90_selftest.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_tp_sm90_selftest.cu new file mode 100644 index 00000000..ccfadbe2 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_tp_sm90_selftest.cu @@ -0,0 +1,427 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// Internal multi-PE self test for the tensor-parallel forward path. +// +// Compiled out unless LIGER_FSLCE_FORWARD_TP_SM90_SELFTEST is defined, so the +// core library's recursive .cu glob picks this file up as an empty translation +// unit and no shared CMake/test file has to change. Build and run with: +// +// nvcc -std=c++17 -arch=sm_90a -O3 --expt-relaxed-constexpr \ +// -c fused_linear_scaled_cross_entropy_forward.cu -o forward.o +// nvcc -std=c++17 -arch=sm_90a -O3 --expt-relaxed-constexpr -rdc=true \ +// -c forward_reduce.cu -o forward_reduce.o +// nvcc -std=c++17 -arch=sm_90a -O3 --expt-relaxed-constexpr -rdc=true \ +// -DLIGER_FSLCE_FORWARD_TP_SM90_SELFTEST \ +// forward_tp_sm90_selftest.cu forward_reduce.o forward.o \ +// -I -I -I \ +// -I$CUTLASS_HOME/include -I$NVSHMEM_HOME/include \ +// -L$NVSHMEM_HOME/lib -lnvshmem_host -lnvshmem_device -lcuda -o tp_selftest +// nvshmrun -n 4 ./tp_selftest +// +// Every PE owns a contiguous shard of one global vocabulary. Each PE rebuilds +// the *whole* weight matrix from the shared seed and checks the reduced +// NLL / LSE / entropy against a full-vocabulary CPU reference, so a wrong +// global max, a missed rescale or a dropped target contribution all show up. +// ═══════════════════════════════════════════════════════════════════════════ + +#if defined(LIGER_FSLCE_FORWARD_TP_SM90_SELFTEST) + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "buffer_pool.cuh" +#include "forward_gemm_sm90.cuh" +#include "forward_reduce.cuh" + +namespace { + +namespace fslce = liger::fused_scaled_linear_cross_entropy; + +using Element = __nv_bfloat16; + +int g_pe = 0; +int g_n_pes = 1; +int g_failures = 0; +ncclComm_t g_nccl_comm = nullptr; + +#define CUDA_OK(expr) \ + do { \ + cudaError_t status = (expr); \ + if (status != cudaSuccess) { \ + std::printf("[pe %d] CUDA error %s at %s:%d: %s\n", \ + g_pe, #expr, __FILE__, __LINE__, \ + cudaGetErrorString(status)); \ + std::exit(1); \ + } \ + } while (0) + +#define NCCL_OK(expr) \ + do { \ + ncclResult_t status = (expr); \ + if (status != ncclSuccess) { \ + std::printf("[pe %d] NCCL error %s at %s:%d: %s\n", \ + g_pe, #expr, __FILE__, __LINE__, ncclGetErrorString(status)); \ + std::exit(1); \ + } \ + } while (0) + +ncclComm_t bootstrap_nccl() { + ncclUniqueId id; + std::memset(&id, 0, sizeof(id)); + if (g_pe == 0) NCCL_OK(ncclGetUniqueId(&id)); + + void* source = nvshmem_malloc(sizeof(id)); + void* destination = nvshmem_malloc(sizeof(id)); + CUDA_OK(cudaMemcpy(source, &id, sizeof(id), cudaMemcpyHostToDevice)); + nvshmem_barrier_all(); + nvshmem_broadcastmem( + NVSHMEM_TEAM_WORLD, + destination, + source, + sizeof(id), + 0); + nvshmem_barrier_all(); + CUDA_OK(cudaMemcpy( + &id, destination, sizeof(id), cudaMemcpyDeviceToHost)); + nvshmem_free(destination); + nvshmem_free(source); + + ncclComm_t comm = nullptr; + NCCL_OK(ncclCommInitRank(&comm, g_n_pes, id, g_pe)); + return comm; +} + +float next_random(std::uint64_t& state) { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + std::uint32_t bits = static_cast(state >> 33); + return static_cast(bits % 2048) / 1024.0f - 1.0f; +} + +// Deterministic from (seed, index) so every PE reconstructs the same global +// tensors without any communication. +float element_at(std::uint64_t seed, std::size_t index, float scale) { + std::uint64_t state = seed ^ (index * 0x9E3779B97F4A7C15ULL); + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + state ^= state >> 29; + state = state * 0xBF58476D1CE4E5B9ULL; + state ^= state >> 32; + std::uint32_t bits = static_cast(state >> 33); + return (static_cast(bits % 2048) / 1024.0f - 1.0f) * scale; +} + +struct ReferenceRow { + float lse; + float nll; + float entropy; + float row_max; + float sum; + float expectation; +}; + +// Full-vocabulary CPU reference, reproducing the kernel's FP16 staging round. +std::vector reference( + const std::vector& x, + const std::vector& weight_global, + const std::vector& target, + int tokens, + int hidden, + int global_vocab, + std::int64_t ignore_index, + float inverse_temperature) { + std::vector rows(tokens); + std::vector scaled(global_vocab); + for (int m = 0; m < tokens; ++m) { + float row_max = -1.0e38f; + for (int v = 0; v < global_vocab; ++v) { + float dot = 0.0f; + for (int k = 0; k < hidden; ++k) { + dot += x[static_cast(m) * hidden + k] * + weight_global[static_cast(v) * hidden + k]; + } + scaled[v] = __half2float(__float2half(dot)) * inverse_temperature; + row_max = fmaxf(row_max, scaled[v]); + } + + double sum = 0.0; + double weighted = 0.0; + for (int v = 0; v < global_vocab; ++v) { + double e = std::exp( + static_cast(scaled[v]) - static_cast(row_max)); + sum += e; + weighted += e * static_cast(scaled[v]); + } + float lse = row_max + static_cast(std::log(sum)); + + bool ignored = target[m] == ignore_index; + rows[m].expectation = static_cast(weighted / sum); + rows[m].row_max = row_max; + rows[m].sum = static_cast(sum); + rows[m].lse = lse; + rows[m].nll = ignored ? 0.0f : lse - scaled[target[m]]; + rows[m].entropy = ignored + ? 0.0f + : lse - static_cast(weighted / sum); + } + return rows; +} + +void expect_near( + const char* label, int row, float actual, float expected, + float absolute, float relative) { + float tolerance = absolute + relative * fabsf(expected); + if (!(fabsf(actual - expected) <= tolerance)) { + if (g_failures < 20) { + std::printf( + "[pe %d] FAIL %-8s row %4d: got %14.6f want %14.6f " + "(tol %.6f)\n", + g_pe, label, row, actual, expected, tolerance); + } + ++g_failures; + } +} + +template +void run_case( + const std::string& name, + int tokens, + int hidden, + int local_vocab, + float temperature, + std::uint64_t seed) { + const int global_vocab = local_vocab * g_n_pes; + const std::int64_t vocab_start = + static_cast(g_pe) * local_vocab; + const float scale = 4.0f / std::sqrt(static_cast(hidden)); + + // Rebuild the global tensors identically on every PE. + std::vector host_x(static_cast(tokens) * hidden); + for (std::size_t i = 0; i < host_x.size(); ++i) { + host_x[i] = element_at(seed, i, scale); + } + std::vector host_w( + static_cast(global_vocab) * hidden); + for (std::size_t i = 0; i < host_w.size(); ++i) { + host_w[i] = element_at(seed ^ 0xABCDEF01ULL, i, scale); + } + // Round through BF16 so host and device see identical operands. + for (auto& v : host_x) v = __bfloat162float(__float2bfloat16(v)); + for (auto& v : host_w) v = __bfloat162float(__float2bfloat16(v)); + + std::vector host_target(tokens); + std::uint64_t target_state = seed ^ 0x5DEECE66DULL; + for (int m = 0; m < tokens; ++m) { + if (m % 13 == 5) { + host_target[m] = -100; // ignore_index + } else { + next_random(target_state); + host_target[m] = static_cast( + (target_state >> 17) % static_cast(global_vocab)); + } + } + + // Device operands: X is replicated, W is this PE's contiguous shard. + std::vector device_x_host(host_x.size()); + for (std::size_t i = 0; i < host_x.size(); ++i) { + device_x_host[i] = __float2bfloat16(host_x[i]); + } + std::vector device_w_host( + static_cast(local_vocab) * hidden); + for (std::size_t i = 0; i < device_w_host.size(); ++i) { + device_w_host[i] = __float2bfloat16( + host_w[static_cast(vocab_start) * hidden + i]); + } + + Element* device_x = nullptr; + Element* device_w = nullptr; + std::int64_t* device_target = nullptr; + CUDA_OK(cudaMalloc(&device_x, device_x_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_w, device_w_host.size() * sizeof(Element))); + CUDA_OK(cudaMalloc(&device_target, tokens * sizeof(std::int64_t))); + CUDA_OK(cudaMemcpy(device_x, device_x_host.data(), + device_x_host.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_w, device_w_host.data(), + device_w_host.size() * sizeof(Element), cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemcpy(device_target, host_target.data(), + tokens * sizeof(std::int64_t), cudaMemcpyHostToDevice)); + + float* device_out = nullptr; + CUDA_OK(cudaMalloc(&device_out, 3 * tokens * sizeof(float))); + CUDA_OK(cudaMemset(device_out, 0xff, 3 * tokens * sizeof(float))); + + fslce::ForwardTpParamsSm90<90> params; + params.gemm.x = device_x; + params.gemm.weight = device_w; + params.gemm.target = device_target; + params.gemm.tokens = tokens; + params.gemm.hidden = hidden; + params.gemm.local_vocab = local_vocab; + params.gemm.vocab_start = vocab_start; + params.gemm.ignore_index = -100; + params.gemm.inverse_temperature = 1.0f / temperature; + params.nll = device_out + 0 * tokens; + params.lse = device_out + 1 * tokens; + params.entropy = device_out + 2 * tokens; + params.nccl_comm_handle = + reinterpret_cast(g_nccl_comm); + + fslce::fused_linear_scaled_cross_entropy_forward< + ReturnEntropy, 90>(params, nullptr); + CUDA_OK(cudaGetLastError()); + CUDA_OK(cudaDeviceSynchronize()); + + std::vector host_out(3 * tokens); + CUDA_OK(cudaMemcpy(host_out.data(), device_out, + host_out.size() * sizeof(float), cudaMemcpyDeviceToHost)); + + fslce::ForwardTpWorkspace probe = + fslce::reserve_forward_tp_workspace(tokens); + std::vector host_global_max(tokens); + std::vector host_reduced( + static_cast(probe.fields) * tokens); + CUDA_OK(cudaMemcpy(host_global_max.data(), probe.global_max, + tokens * sizeof(float), cudaMemcpyDeviceToHost)); + CUDA_OK(cudaMemcpy(host_reduced.data(), probe.reduced, + host_reduced.size() * sizeof(float), cudaMemcpyDeviceToHost)); + + std::vector expected = reference( + host_x, host_w, host_target, tokens, hidden, global_vocab, + params.gemm.ignore_index, params.gemm.inverse_temperature); + + int before = g_failures; + int stage_max_bad = 0; + int stage_sum_bad = 0; + for (int m = 0; m < tokens; ++m) { + if (fabsf(host_global_max[m] - expected[m].row_max) > + 2.0e-3f + 1.0e-3f * fabsf(expected[m].row_max)) { + ++stage_max_bad; + } + float reduced_sum = host_reduced[0 * tokens + m]; + if (fabsf(reduced_sum - expected[m].sum) > + 1.0e-2f + 5.0e-3f * fabsf(expected[m].sum)) { + ++stage_sum_bad; + } + } + if (stage_max_bad || stage_sum_bad) { + std::printf("[pe %d] STAGE %-28s max_bad=%d sum_bad=%d of %d rows\n", + g_pe, name.c_str(), stage_max_bad, stage_sum_bad, tokens); + } + for (int m = 0; m < tokens; ++m) { + expect_near("lse", m, host_out[1 * tokens + m], + expected[m].lse, 3.0e-3f, 2.0e-3f); + expect_near("nll", m, host_out[0 * tokens + m], + expected[m].nll, 3.0e-3f, 2.0e-3f); + if constexpr (ReturnEntropy) { + expect_near("entropy", m, host_out[2 * tokens + m], + expected[m].entropy, 5.0e-3f, 5.0e-3f); + // Backward recovers the softmax-weighted mean logit from the two + // saved output tensors as (lse - entropy); check that identity + // holds wherever the gradient is non-zero. + if (host_target[m] != params.gemm.ignore_index) { + expect_near("lse-entropy", m, + host_out[1 * tokens + m] - host_out[2 * tokens + m], + expected[m].expectation, 5.0e-3f, 5.0e-3f); + } + } + if (g_failures > before + 20) break; + } + + if (g_pe == 0) { + std::printf("%-40s %s (PEs=%d M=%d H=%d Vloc=%d Vglobal=%d T=%.2f " + "entropy=%d)\n", + name.c_str(), + g_failures == before ? "PASS" : "FAIL", + g_n_pes, tokens, hidden, local_vocab, global_vocab, temperature, + static_cast(ReturnEntropy)); + } + + // Keep every PE on the same case boundary before pooled symmetric + // reduction buffers are reused by the next launch. + nvshmem_barrier_all(); + + CUDA_OK(cudaFree(device_x)); + CUDA_OK(cudaFree(device_w)); + CUDA_OK(cudaFree(device_target)); + CUDA_OK(cudaFree(device_out)); +} + +} // namespace + +int main() { + nvshmem_init(); + g_pe = nvshmem_my_pe(); + g_n_pes = nvshmem_n_pes(); + CUDA_OK(cudaSetDevice(nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE))); + + cudaDeviceProp properties = {}; + CUDA_OK(cudaGetDeviceProperties(&properties, nvshmem_team_my_pe( + NVSHMEMX_TEAM_NODE))); + if (properties.major != 9) { + if (g_pe == 0) { + std::printf("skipping: requires Hopper (sm90), found sm_%d%d\n", + properties.major, properties.minor); + } + nvshmem_finalize(); + return 0; + } + + if (g_pe == 0) { + std::printf("tensor-parallel forward over %d PEs\n", g_n_pes); + } + g_nccl_comm = bootstrap_nccl(); + + // Capacities are immutable once allocated, so configure for the largest + // shape used below before any launch. + constexpr int kMaxTokens = 1024; + constexpr int kMaxLocalVocab = 2560; + fslce::configure_forward_tp_workspace(kMaxTokens, kMaxLocalVocab); + + try { + run_case("tp aligned / no entropy", 256, 128, 640, 1.0f, 101); + run_case("tp aligned / entropy", 256, 128, 640, 1.0f, 102); + run_case("tp ragged M/H/V", 200, 104, 517, 0.8f, 103); + run_case("tp temperature / no entropy", 384, 192, 901, 1.4f, 104); + run_case("tp large split", 1024, 128, 2560, 1.0f, 105); + } catch (const std::exception& error) { + std::printf("[pe %d] exception: %s\n", g_pe, error.what()); + ncclCommDestroy(g_nccl_comm); + nvshmem_finalize(); + return 1; + } + + // Make every rank's verdict visible before PE 0 prints the summary. + // Pooled, not raw nvshmem_malloc: this module never runs its own symmetric + // allocation lifecycle. + auto& pool = liger::global_buffer_pool(); + int* local_slot = static_cast( + pool.get_symmetric("fslce_tp_selftest_verdict_local", sizeof(int))); + int* total_slot = static_cast( + pool.get_symmetric("fslce_tp_selftest_verdict_total", sizeof(int))); + CUDA_OK(cudaMemcpy(local_slot, &g_failures, sizeof(int), + cudaMemcpyHostToDevice)); + nvshmem_int_sum_reduce(NVSHMEM_TEAM_WORLD, total_slot, local_slot, 1); + int host_total = 0; + CUDA_OK(cudaMemcpy(&host_total, total_slot, sizeof(int), + cudaMemcpyDeviceToHost)); + + if (g_pe == 0) { + std::printf("%s: %d comparison failures across %d PEs\n", + host_total == 0 ? "OK" : "FAILED", host_total, g_n_pes); + } + NCCL_OK(ncclCommDestroy(g_nccl_comm)); + fslce::reset_forward_tp_workspace_configuration(); + liger::global_buffer_pool().clear(); + nvshmem_finalize(); + return host_total == 0 ? 0 : 1; +} + +#endif // LIGER_FSLCE_FORWARD_TP_SM90_SELFTEST diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu new file mode 100644 index 00000000..6845c00d --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu @@ -0,0 +1,542 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy backward host dispatch. dZ always +// selects the zero-spill cluster-2 handoff. NVLS and hierarchical paths use +// the four-stage combined dX+dW cluster kernel; direct-peer uses its fallback. +// ═══════════════════════════════════════════════════════════════════════════ + +// Must precede every other CUTLASS/CuTe include: it opts the TU into CuTe's +// extended SM90 MMA shapes. +#include "backward_gemm_mainloop_sm90.cuh" + +#include "backward_gemm_sm90.cuh" +#include "backward_reduce_sm90.cuh" +#include "gemm_sm90.cuh" +#include "liger_cute/check.h" +#include "liger_cute/detail/tp_reduce.cuh" +#include "workspace.cuh" + +#include +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { +namespace { + +using namespace cute; + +using HostConfig = BackwardGemmConfigSm90<90>; +using HostLaunch = BackwardGemmLaunchSm90<90>; +using HostTraits = BackwardGemmTraitsSm90<90>; +using HostCombinedDwTraits = BackwardGemmTraitsSm90<90, 4, 64, 2>; +using HostDxClusterTraits = BackwardDxCluster2TraitsSm90<90>; +using HostDzHandoffTraits = BackwardDzHandoffTraitsSm90<90>; +using HostElement = HostTraits::Element; + +// The staging ring depth comes from dx_reduce.cuh so the kernel template, +// symmetric capacity and CTA-local mbarrier arrays cannot drift apart. + +void check_cuda(cudaError_t error, const char* what) { + LIGER_CHECK( + error == cudaSuccess, + "fused_scaled_linear_cross_entropy backward: ", + what, + " failed: ", + cudaGetErrorString(error)); +} + +template +auto make_row_major_load(const Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_LOAD{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_row_major_load_tile( + const Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_LOAD{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}, + make_shape(Int{}, Int{}), + _1{}); +} + +template +auto make_row_major_store(Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_STORE{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_row_major_reduce_add(Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}); +} + +template +auto make_row_major_store_tile(Element* ptr, int64_t rows, int64_t cols) { + return make_tma_copy( + SM90_TMA_STORE{}, + backward_row_major_tensor(ptr, rows, cols), + Layout{}, + make_shape(Int{}, Int{}), + _1{}); +} + +template +auto make_col_major_load( + const Element* ptr, int64_t rows, int64_t cols, int64_t leading) { + return make_tma_copy( + SM90_TMA_LOAD{}, + backward_col_major_tensor(ptr, rows, cols, leading), + Layout{}); +} + +template < + class Layout, + int TileRows, + int TileCols, + int Multicast, + class Element> +auto make_col_major_load_multicast( + const Element* ptr, int64_t rows, int64_t cols, int64_t leading) { + return make_tma_copy( + SM90_TMA_LOAD_MULTICAST{}, + backward_col_major_tensor(ptr, rows, cols, leading), + Layout{}, + make_shape(Int{}, Int{}), + Int{}); +} + +void validate(const BackwardTpParamsSm90<90>& params) { + const auto& gemm = params.gemm; + LIGER_CHECK(gemm.tokens >= 0, "tokens must be non-negative"); + LIGER_CHECK(gemm.hidden > 0, "hidden must be positive"); + LIGER_CHECK(gemm.local_vocab > 0, "local_vocab must be positive"); + LIGER_CHECK( + gemm.x != nullptr && gemm.weight != nullptr && + gemm.target != nullptr && gemm.grad_output != nullptr && + gemm.lse != nullptr, + "backward inputs x / weight / target / grad_output / lse must be " + "non-null"); + LIGER_CHECK( + gemm.grad_input != nullptr && gemm.grad_weight != nullptr, + "backward outputs grad_input and grad_weight must be non-null"); + LIGER_CHECK( + gemm.dz_workspace != nullptr, + "the dZ wave workspace must be reserved"); + LIGER_CHECK( + params.tiles_per_reduce == 1 || params.tiles_per_reduce == 2 || + params.tiles_per_reduce == 4, + "TilesPerReduce must be 1, 2 or 4, got ", + params.tiles_per_reduce); + LIGER_CHECK( + params.num_comm_channels >= 1, + "num_comm_channels must be positive (legacy compatibility argument)"); + + // TMA needs 16 B aligned base addresses and row strides for its BF16 + // operands. dX uses generic 8 B vector stores. + auto aligned = [](const void* pointer) { + return reinterpret_cast(pointer) % 16 == 0; + }; + LIGER_CHECK( + aligned(gemm.x) && aligned(gemm.weight) && + aligned(gemm.grad_weight) && aligned(gemm.dz_workspace), + "backward TMA BF16 operands must be 16 B aligned"); + LIGER_CHECK( + reinterpret_cast(gemm.grad_input) % 8 == 0, + "grad_input must be 8 B aligned for vectorized BF16 stores"); + LIGER_CHECK( + (static_cast(gemm.hidden) * sizeof(HostElement)) % 16 == 0, + "hidden * sizeof(bfloat16) must be a multiple of 16 B for TMA"); + validate_backward_tp_shape( + gemm.tokens, gemm.hidden, gemm.local_vocab); + + std::size_t needed = HostLaunch::dz_workspace_bytes(gemm.local_vocab); + LIGER_CHECK( + gemm.dz_workspace_bytes >= needed, + "the dZ wave workspace needs ", + needed, + " B, got ", + gemm.dz_workspace_bytes); +} + +template < + int TilesPerReduce, + bool ReturnEntropy, + liger_cute::detail::LocalReduceBackend Backend, + bool RequiresRemote> +void launch_instance( + const BackwardTpParamsSm90<90>& params, + const liger_cute::detail::TpReducePlan& reduce, + cudaStream_t stream) { + static constexpr bool kUseNvls = + Backend == liger_cute::detail::LocalReduceBackend::kNvls; + static_assert( + !RequiresRemote || kUseNvls, + "the remote follow-up requires a local NVLS reduce-scatter"); + using CommConfig = + DxCommConfig; + using DzSmem = BackwardDzHandoffSmemSm90<90, ReturnEntropy>; + static constexpr bool kUseCluster = kUseNvls; + + const auto& gemm = params.gemm; + int padded_vocab = HostLaunch::padded_vocab(gemm.local_vocab); + + auto* x = static_cast(gemm.x); + auto* w = static_cast(gemm.weight); + auto* dz = static_cast(gemm.dz_workspace); + auto* dw_out = static_cast(gemm.grad_weight); + DxReduceWorkspace comm = {}; + + int resident_ctas = backward_dx_resident_cta_capacity(); + constexpr int kClusterStoreRows = + kMaxDxResidentCtas * kDxCommWarpsPerChannel * + kDxRingStages * TilesPerReduce * HostDxClusterTraits::kTileM; + int cluster_store_rows = + resident_ctas * kDxCommWarpsPerChannel * + kDxRingStages * TilesPerReduce * HostDxClusterTraits::kTileM; + using ClusterStoreTma = decltype( + make_row_major_store< + typename HostDxClusterTraits::SmemStore::Single>( + static_cast(nullptr), + kClusterStoreRows, + HostDxClusterTraits::kTileN)); + ClusterStoreTma cluster_tma_store = {}; + if constexpr (kUseCluster) { + LIGER_CHECK( + reduce.backend == + liger_cute::detail::LocalReduceBackend::kNvls, + "fused_scaled_linear_cross_entropy backward: clustered dX " + "requires a node-local NVLS team"); + LIGER_CHECK( + reduce.remote.enabled() == RequiresRemote, + "fused_scaled_linear_cross_entropy backward: remote reduction " + "dispatch does not match the configured plan"); + comm = reserve_dx_reduce_workspace( + TilesPerReduce, + kDxRingStages, + resident_ctas); + cluster_tma_store = + make_row_major_store< + typename HostDxClusterTraits::SmemStore::Single>( + comm.partial, + cluster_store_rows, + HostDxClusterTraits::kTileN); + } + auto cluster_tma_dz = + make_row_major_load( + dz, + HostConfig::kWaveRows, + padded_vocab); + auto cluster_tma_wt = make_col_major_load_multicast< + typename HostDxClusterTraits::SmemB::Single, + HostDxClusterTraits::kFragmentTileN, + HostDxClusterTraits::kTileK, + HostDxClusterTraits::kClusterM>( + w, + gemm.hidden, + gemm.local_vocab, + gemm.hidden); + auto cluster_tma_dzt = make_col_major_load_multicast< + typename HostCombinedDwTraits::SmemAMnMajor::Single, + HostCombinedDwTraits::kTileM, + HostCombinedDwTraits::kTileK, + HostDxClusterTraits::kClusterM>( + dz, + padded_vocab, + HostConfig::kWaveRows, + padded_vocab); + auto cluster_tma_xt = make_col_major_load< + typename HostCombinedDwTraits::SmemBMnMajor::Single>( + x, gemm.hidden, gemm.tokens, gemm.hidden); + auto cluster_tma_dw_store = make_row_major_store< + typename HostCombinedDwTraits::DwSmemStore1::Single>( + dw_out, gemm.local_vocab, gemm.hidden); + auto cluster_tma_dw_add = make_row_major_reduce_add< + typename HostCombinedDwTraits::DwSmemStore1::Single>( + dw_out, gemm.local_vocab, gemm.hidden); + using ClusterBundle = BackwardClusterCombinedTmaBundleSm90< + decltype(cluster_tma_dz), + decltype(cluster_tma_wt), + ClusterStoreTma, + decltype(cluster_tma_dzt), + decltype(cluster_tma_xt), + decltype(cluster_tma_dw_store), + decltype(cluster_tma_dw_add)>; + ClusterBundle cluster_bundle{ + cluster_tma_dz, + cluster_tma_wt, + cluster_tma_store, + cluster_tma_dzt, + cluster_tma_xt, + cluster_tma_dw_store, + cluster_tma_dw_add}; + const bool use_split_k2 = + backward_dx_split_k(gemm.hidden, gemm.local_vocab) == 2; + auto* cluster_kernel = use_split_k2 + ? &backward_dx_cluster2_gemm_wave_kernel_sm90< + 2, 90, CommConfig, ClusterBundle> + : &backward_dx_cluster2_gemm_wave_kernel_sm90< + 1, 90, CommConfig, ClusterBundle>; + + auto tma_x = + make_row_major_load_tile< + typename HostDzHandoffTraits::SmemTraitsA::Single, + HostDzHandoffTraits::kTileM, + HostDzHandoffTraits::kTileK>( + x, gemm.tokens, gemm.hidden); + auto tma_w = + make_row_major_load_tile< + typename HostDzHandoffTraits::SmemTraitsB::Single, + HostDzHandoffTraits::kTileN, + HostDzHandoffTraits::kTileK>( + w, gemm.local_vocab, gemm.hidden); + auto tma_dz_store = + make_row_major_store_tile< + typename HostDzHandoffTraits::SmemStoreSingle, + HostDzHandoffTraits::kTileM, + HostDzHandoffTraits::kStoreTileN>( + dz, HostConfig::kWaveRows, padded_vocab); + auto* dz_kernel_fn = &backward_dz_handoff_wave_kernel_sm90< + ReturnEntropy, + 90, + decltype(tma_x), + decltype(tma_w), + decltype(tma_dz_store)>; + constexpr int kDzSmemBytes = static_cast(sizeof(DzSmem)); + static_assert(kDzSmemBytes <= kBackwardMaxSmemBytes, + "SM90 backward dZ shared-memory footprint exceeds the Hopper limit"); + using DzCluster = sm90::ClusterLaunchSm90<90>; + check_cuda( + DzCluster::prepare(dz_kernel_fn, kDzSmemBytes), + "cudaFuncSetAttribute(dZ MaxDynamicSharedMemorySize / " + "NonPortableClusterSizeAllowed)"); + using DxCluster = sm90::ClusterLaunchSm90<90>; + constexpr int kClusterSmemBytes = + static_cast(sizeof(BackwardDxCluster2SmemSm90<90>)); + static_assert(kClusterSmemBytes == 230400, + "production combined dX+dW shared memory must remain 230400 B"); + static_assert(kClusterSmemBytes <= kBackwardMaxSmemBytes, + "SM90 dX cluster2 shared-memory footprint exceeds Hopper limit"); + static_assert(HostConfig::kMTilesPerWave % + HostDxClusterTraits::kClusterM == 0); + int dx_cluster_count = 0; + if constexpr (kUseCluster) { + check_cuda( + DxCluster::prepare(cluster_kernel, kClusterSmemBytes), + "cudaFuncSetAttribute(dX cluster2 MaxDynamicSharedMemorySize / " + "NonPortableClusterSizeAllowed)"); + int max_active_clusters = DxCluster::max_active_clusters( + cluster_kernel, + HostDxClusterTraits::kNumThreads, + kClusterSmemBytes, + HostDxClusterTraits::kClusterM); + LIGER_CHECK( + max_active_clusters > 0, + "fused_scaled_linear_cross_entropy backward: " + "no active dX cluster2 fits"); + int total_cluster_tiles = + (HostConfig::kMTilesPerWave / + HostDxClusterTraits::kClusterM) * + HostLaunch::num_dx_n_tiles(gemm.hidden) * + (use_split_k2 ? 2 : 1); + int dw_tiles = HostLaunch::num_dw_m_tiles(gemm.local_vocab) * + HostLaunch::num_dw_n_tiles(gemm.hidden); + int combined_tiles = total_cluster_tiles > ceil_div(dw_tiles, 2) + ? total_cluster_tiles + : ceil_div(dw_tiles, 2); + dx_cluster_count = + combined_tiles < max_active_clusters + ? combined_tiles + : max_active_clusters; + LIGER_CHECK( + dx_cluster_count * HostDxClusterTraits::kClusterM <= + resident_ctas, + "fused_scaled_linear_cross_entropy backward: dX cluster grid " + "exceeds the configured resident CTA capacity"); + } + + BackwardGemmParamsSm90<90> device_params = gemm; + + int dz_work = HostConfig::kMTilesPerWave * ceil_div( + padded_vocab, + HostDzHandoffTraits::kClusterM * + HostDzHandoffTraits::kTileN); + int max_dz_clusters = DzCluster::max_active_clusters( + dz_kernel_fn, + HostConfig::kNumThreads, + kDzSmemBytes, + HostDzHandoffTraits::kClusterM); + LIGER_CHECK( + max_dz_clusters > 0, + "fused_scaled_linear_cross_entropy backward: " + "no active dZ handoff clusters fit"); + int dz_clusters = + dz_work < max_dz_clusters ? dz_work : max_dz_clusters; + + int dx_groups_per_wave = HostConfig::kMTilesPerWave * + ceil_div(HostLaunch::num_dx_n_tiles(gemm.hidden), TilesPerReduce); + int dw_total = HostLaunch::num_dw_m_tiles(gemm.local_vocab) * + HostLaunch::num_dw_n_tiles(gemm.hidden); + int work = dx_groups_per_wave; + if (dw_total > work) work = dw_total; + int grid = work < resident_ctas ? work : resident_ctas; + int num_waves = HostLaunch::num_waves(gemm.tokens); + if constexpr (!kUseCluster) { + comm = reserve_dx_reduce_workspace( + TilesPerReduce, kDxRingStages, grid); + } + liger_cute::detail::begin_tp_reduce(comm.launch_epoch, stream); + if constexpr (kUseCluster) { + if (use_split_k2) { + std::size_t durable_elements = + static_cast(num_waves) * + HostConfig::kWaveRows * + static_cast( + HostLaunch::num_dx_n_tiles(gemm.hidden)) * + HostDxClusterTraits::kTileN / + static_cast( + reduce.nvls.size); + check_cuda( + cudaMemsetAsync( + reduce.nvls.reduced_shard, + 0, + durable_elements * sizeof(float), + stream), + "clear split-K dX durable output"); + } + } + for (int wave = 0; wave < num_waves; ++wave) { + check_cuda( + DzCluster::launch( + dz_kernel_fn, + dim3( + HostDzHandoffTraits::kClusterM, + 1, + static_cast(dz_clusters)), + HostConfig::kNumThreads, + kDzSmemBytes, + HostDzHandoffTraits::kClusterM, + stream, + tma_x, + tma_w, + tma_dz_store, + device_params, + wave), + "cudaLaunchKernelEx(backward_dz_handoff_wave_kernel_sm90)"); + if constexpr (kUseCluster) { + check_cuda( + DxCluster::launch( + cluster_kernel, + dim3( + HostDxClusterTraits::kClusterM, + 1, + static_cast(dx_cluster_count)), + HostDxClusterTraits::kNumThreads, + kClusterSmemBytes, + HostDxClusterTraits::kClusterM, + stream, + cluster_bundle, + comm, + reduce.nvls, + device_params, + wave), + "cudaLaunchKernelEx(" + "backward_dx_cluster2_gemm_wave_kernel_sm90)"); + } else { + launch_backward_dx_dw_direct_peer_wave_sm90( + ReturnEntropy, + TilesPerReduce, + device_params, + comm, + reduce.direct, + grid, + stream, + wave); + } + } + if constexpr (kUseCluster) { + int cluster_grid = + dx_cluster_count * HostDxClusterTraits::kClusterM; + finalize_backward_dx_cluster_sm90( + device_params, + comm, + reduce.nvls, + reduce.remote, + cluster_grid, + num_waves, + stream); + } + liger_cute::detail::end_tp_reduce(stream); +} + +template +void dispatch_instance( + const BackwardTpParamsSm90<90>& params, + const liger_cute::detail::TpReducePlan& reduce, + cudaStream_t stream) { + if (reduce.backend == + liger_cute::detail::LocalReduceBackend::kNvls) { + if (reduce.remote.enabled()) { + launch_instance< + TilesPerReduce, + ReturnEntropy, + liger_cute::detail::LocalReduceBackend::kNvls, + true>(params, reduce, stream); + } else { + launch_instance< + TilesPerReduce, + ReturnEntropy, + liger_cute::detail::LocalReduceBackend::kNvls, + false>(params, reduce, stream); + } + } else { + launch_instance< + TilesPerReduce, + ReturnEntropy, + liger_cute::detail::LocalReduceBackend::kDirectPeer, + false>(params, reduce, stream); + } +} + +} // namespace + +template +void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90& params, cudaStream_t stream) { + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + if (params.gemm.tokens == 0) return; + validate(params); + + liger_cute::detail::TpReducePlan reduce = + liger_cute::detail::tp_reduce_plan(); + switch (params.tiles_per_reduce) { + case 1: + dispatch_instance<1, ReturnEntropy>(params, reduce, stream); + break; + case 2: + dispatch_instance<2, ReturnEntropy>(params, reduce, stream); + break; + default: + dispatch_instance<4, ReturnEntropy>(params, reduce, stream); + break; + } +} + +template void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90<90>&, + cudaStream_t); +template void fused_linear_scaled_cross_entropy_backward( + const BackwardTpParamsSm90<90>&, + cudaStream_t); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu new file mode 100644 index 00000000..3ee3ffec --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu @@ -0,0 +1,348 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 tensor-parallel fused scaled linear cross entropy — forward launcher. +// +// Sole production host call graph: +// +// fused_linear_scaled_cross_entropy_forward +// -> forward_gemm_kernel_sm90 +// -> forward_split_reduce_kernel_sm90 +// -> NCCL MAX +// -> pack_forward_reduction_kernel +// -> NCCL SUM +// -> finalize_forward_kernel +// +// The WGMMA translation unit is deliberately non-RDC. NVSHMEM collective +// kernels remain in forward_reduce.cu, which is compiled with RDC. +// ═══════════════════════════════════════════════════════════════════════════ + +// Must precede every other CUTLASS/CuTe include: it opts the TU into CuTe's +// extended SM90 MMA shapes, which is where the N160 WGMMA atom lives. +#include "forward_gemm_mainloop_sm90.cuh" + +#include "forward_gemm_sm90.cuh" +#include "forward_reduce.cuh" +#include "gemm_sm90.cuh" +#include "liger_cute/check.h" +#include "online_softmax.cuh" + +#include + +#include +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { +namespace { + +template +__global__ void pack_forward_reduction_kernel( + const float* local_max, + const float* global_max, + const float* local_sum, + const float* local_target, + const float* local_weighted_sum, + float* packed, + int tokens) { + int row = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (row >= tokens) return; + + float correction = local_sum[row] == 0.0f + ? 0.0f + : expf(local_max[row] - global_max[row]); + packed[kForwardReducedSumField * tokens + row] = + local_sum[row] * correction; + packed[kForwardReducedTargetField * tokens + row] = + local_target[row]; + if constexpr (ReturnEntropy) { + packed[kForwardReducedWeightedField * tokens + row] = + local_weighted_sum[row] * correction; + } +} + +template +__global__ void finalize_forward_kernel( + const float* global_max, + const float* reduced, + const std::int64_t* target, + float* nll, + float* lse, + float* entropy, + int tokens, + std::int64_t ignore_index) { + int row = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (row >= tokens) return; + + float global_weighted = 0.0f; + if constexpr (ReturnEntropy) { + global_weighted = + reduced[kForwardReducedWeightedField * tokens + row]; + } + FinalizedSoftmax result = finalize_softmax( + global_max[row], + reduced[kForwardReducedSumField * tokens + row], + reduced[kForwardReducedTargetField * tokens + row], + global_weighted, + target[row] == ignore_index); + nll[row] = result.nll; + lse[row] = result.lse; + if constexpr (ReturnEntropy) { + entropy[row] = result.entropy; + } +} + +void check_cuda(cudaError_t error, const char* what) { + LIGER_CHECK( + error == cudaSuccess, + "fused_scaled_linear_cross_entropy forward: ", + what, + " failed: ", + cudaGetErrorString(error)); +} + +void check_nccl(ncclResult_t result, const char* what) { + LIGER_CHECK( + result == ncclSuccess, + "fused_scaled_linear_cross_entropy forward: ", + what, + " failed: ", + ncclGetErrorString(result)); +} + +} // namespace + +template +void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90& params, + cudaStream_t stream) { + using Traits = ForwardGemmTraitsSm90; + using Config = typename Traits::Config; + using Launch = ForwardGemmLaunchSm90; + using Element = typename Traits::Element; + using Smem = ForwardGemmSmemSm90; + + static_assert(Compute == 90, + "SM90 fused scaled linear cross entropy requires Compute=90"); + + const ForwardGemmParamsSm90& input = params.gemm; + LIGER_CHECK(input.tokens >= 0, "tokens must be non-negative"); + if (input.tokens == 0) return; + LIGER_CHECK(input.hidden > 0, "hidden must be positive"); + LIGER_CHECK(input.local_vocab > 0, "local_vocab must be positive"); + LIGER_CHECK( + input.x != nullptr && input.weight != nullptr && + input.target != nullptr, + "forward inputs must be non-null"); + LIGER_CHECK( + params.nll != nullptr && params.lse != nullptr, + "forward TP outputs nll and lse must be non-null"); + if constexpr (ReturnEntropy) { + LIGER_CHECK( + params.entropy != nullptr, + "entropy output is required when ReturnEntropy=true"); + } + LIGER_CHECK( + params.nccl_comm_handle != 0, + "forward TP requires a valid ncclComm_t handle"); + LIGER_CHECK( + reinterpret_cast(input.x) % 16 == 0 && + reinterpret_cast(input.weight) % 16 == 0, + "forward GEMM operands must be 16 B aligned for TMA"); + LIGER_CHECK( + (static_cast(input.hidden) * sizeof(Element)) % 16 == 0, + "hidden * sizeof(bfloat16) must be a multiple of 16 B for TMA"); + LIGER_CHECK( + input.tuning.split_n >= 0 && input.tuning.base_split_n >= 0 && + input.tuning.extra_m_pairs >= 0 && + input.tuning.target_cluster_pairs >= 0, + "split and cluster tuning values must be non-negative"); + LIGER_CHECK(input.tuning.max_split_n >= 1, "max_split_n must be positive"); + LIGER_CHECK( + input.tuning.extra_m_pairs == 0 || + input.tuning.base_split_n != 0, + "extra_m_pairs requires an explicit base_split_n"); + + ForwardTpWorkspace reduce_workspace = + reserve_forward_tp_workspace(input.tokens); + + ForwardGemmParamsSm90 gemm_params = input; + gemm_params.output = reduce_workspace.local; + gemm_params.workspace = reduce_workspace.gemm_split_partials; + gemm_params.workspace_bytes = + reduce_workspace.gemm_split_partials_bytes; + + LIGER_CHECK( + gemm_params.output.local_max != nullptr && + gemm_params.output.local_sum != nullptr && + gemm_params.output.local_target != nullptr, + "forward GEMM local statistics buffers must be non-null"); + if constexpr (ReturnEntropy) { + LIGER_CHECK( + gemm_params.output.local_weighted_sum != nullptr, + "local_weighted_sum is required when ReturnEntropy=true"); + } + + // Exact CuTe-DSL operand ordering: X is an ordinary TMA load and each + // N160 weight panel is multicast over the cluster's M mode. + auto tma_load_x = sm90::TmaSm90::template make_load< + typename Traits::SmemLayoutX1, + Traits::kTileM, + Traits::kTileK>( + static_cast(gemm_params.x), + gemm_params.tokens, + gemm_params.hidden); + auto tma_load_w = + sm90::TmaSm90::template make_load_multicast< + typename Traits::SmemLayoutW1, + Traits::kAccumulatorTileN, + Traits::kTileK, + Traits::kClusterM>( + static_cast(gemm_params.weight), + gemm_params.local_vocab, + gemm_params.hidden); + + auto* kernel_fn = &forward_gemm_kernel_sm90< + ReturnEntropy, Compute, decltype(tma_load_x), decltype(tma_load_w)>; + + constexpr int kSmemBytes = static_cast(sizeof(Smem)); + static_assert(kSmemBytes <= kHopperMaxSmemBytes, + "SM90 forward shared-memory footprint exceeds the Hopper limit"); + static_assert(kSmemBytes <= Launch::smem_bytes(ReturnEntropy), + "ForwardGemmLaunchSm90::smem_bytes must bound the real footprint"); + + using Cluster = sm90::ClusterLaunchSm90; + check_cuda( + Cluster::prepare(kernel_fn, kSmemBytes), + "cudaFuncSetAttribute(MaxDynamicSharedMemorySize / " + "NonPortableClusterSizeAllowed)"); + + int max_active_cluster_pairs = Cluster::max_active_clusters( + kernel_fn, Config::kNumThreads, kSmemBytes, Config::kClusterM); + ForwardGemmSplitSm90 split = Launch::resolve_split( + gemm_params.tuning, + gemm_params.tokens, + gemm_params.local_vocab, + max_active_cluster_pairs); + LIGER_CHECK( + split.base_split_n >= 1 && + split.split_n <= split.num_logical_n_tiles, + "split_n must be between 1 and the number of logical vocabulary tiles"); + LIGER_CHECK( + split.extra_m_pairs < split.num_m_pairs && + split.split_n == + split.base_split_n + (split.extra_m_pairs != 0 ? 1 : 0), + "uneven split tuning must use split_n=base_split_n+1 for a proper " + "subset of M pairs"); + + // Split partials are laid out [statistic][M tile][split][row]. + std::size_t partial_rows = + static_cast(split.num_m_tiles) * + static_cast(split.split_n) * + static_cast(Config::kTileM); + std::size_t partial_bytes = partial_rows * sizeof(float); + std::size_t required_bytes = + partial_bytes * (ReturnEntropy ? 4u : 3u); + LIGER_CHECK( + gemm_params.workspace != nullptr && + gemm_params.workspace_bytes >= required_bytes, + "forward GEMM needs a ", + required_bytes, + " B split-partial workspace, got ", + gemm_params.workspace_bytes); + LIGER_CHECK( + reinterpret_cast(gemm_params.workspace) % 16 == 0, + "forward GEMM workspace must be 16 B aligned"); + + float* scratch = static_cast(gemm_params.workspace); + ForwardGemmPartialsSm90 partials; + partials.partial_max = scratch + 0 * partial_rows; + partials.partial_sum = scratch + 1 * partial_rows; + partials.partial_target = scratch + 2 * partial_rows; + partials.partial_weighted = + ReturnEntropy ? scratch + 3 * partial_rows : nullptr; + + check_cuda( + Cluster::launch( + kernel_fn, + dim3( + static_cast(Config::kClusterM), + 1u, + static_cast(split.num_cluster_pairs)), + Config::kNumThreads, + kSmemBytes, + Config::kClusterM, + stream, + tma_load_x, + tma_load_w, + gemm_params, + partials, + split), + "cudaLaunchKernelEx(forward_gemm_kernel_sm90)"); + + forward_split_reduce_kernel_sm90 + <<>>( + gemm_params, partials, split); + check_cuda( + cudaGetLastError(), "forward_split_reduce_kernel_sm90 launch"); + + ncclComm_t comm = + reinterpret_cast(params.nccl_comm_handle); + check_nccl( + ncclAllReduce( + reduce_workspace.local.local_max, + reduce_workspace.global_max, + gemm_params.tokens, + ncclFloat32, + ncclMax, + comm, + stream), + "ncclAllReduce(MAX)"); + + constexpr int kPostThreads = 256; + int post_blocks = ceil_div(gemm_params.tokens, kPostThreads); + pack_forward_reduction_kernel + <<>>( + reduce_workspace.local.local_max, + reduce_workspace.global_max, + reduce_workspace.local.local_sum, + reduce_workspace.local.local_target, + reduce_workspace.local.local_weighted_sum, + reduce_workspace.packed, + gemm_params.tokens); + check_cuda( + cudaGetLastError(), "pack_forward_reduction_kernel launch"); + + check_nccl( + ncclAllReduce( + reduce_workspace.packed, + reduce_workspace.reduced, + static_cast(reduce_workspace.fields) * + static_cast(gemm_params.tokens), + ncclFloat32, + ncclSum, + comm, + stream), + "ncclAllReduce(SUM)"); + + finalize_forward_kernel + <<>>( + reduce_workspace.global_max, + reduce_workspace.reduced, + gemm_params.target, + params.nll, + params.lse, + params.entropy, + gemm_params.tokens, + gemm_params.ignore_index); + check_cuda( + cudaGetLastError(), "finalize_forward_kernel launch"); +} + +template void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90<90>&, + cudaStream_t); +template void fused_linear_scaled_cross_entropy_forward( + const ForwardTpParamsSm90<90>&, + cudaStream_t); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_scaled_linear_cross_entropy.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_scaled_linear_cross_entropy.cuh new file mode 100644 index 00000000..f385538b --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_scaled_linear_cross_entropy.cuh @@ -0,0 +1,10 @@ +#pragma once + +#include "backward.cuh" +#include "backward_gemm_sm90.cuh" +#include "config.cuh" +#include "dx_reduce.cuh" +#include "forward_gemm_sm90.cuh" +#include "forward_reduce.cuh" +#include "online_softmax.cuh" +#include "workspace.cuh" diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/gemm_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/gemm_sm90.cuh new file mode 100644 index 00000000..fdc7196f --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/gemm_sm90.cuh @@ -0,0 +1,506 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// Reusable SM90 (Hopper) warp-specialized GEMM mechanics. +// +// Architecture-specific building blocks shared by the fused scaled linear +// cross entropy forward path and the (upcoming) backward dZ/dX/dW paths, so +// neither has to copy/paste TMA descriptor construction, SMEM swizzle atom +// selection, WGMMA atom selection, mbarrier pipeline wiring, warp-group +// register reconfiguration or the cluster launch dance. +// +// Design rules: +// * Zero-overhead: every entity is a template or a CUTE_DEVICE +// (__device__ __forceinline__) function, and the mainloop drivers take +// the caller's issue/fence steps as inlined functors. The generated SASS +// is bit-identical to the hand-written loops these replaced — see the +// cuobjdump diff in the module's validation notes. +// * Nothing here decides *ordering policy*. MainloopKLoopSm90 encodes the +// one-stage release lag the CuTe-DSL kernels rely on and nothing else; +// what to issue per stage stays with the caller. +// * `int Compute` on every GEMM-related template, defaulted to 90, with a +// static_assert so a future SM100 port is a new header, not a silent +// mis-instantiation of this one. +// ═══════════════════════════════════════════════════════════════════════════ + +// The extended SM90 MMA shapes (e.g. the N160 WGMMA atom) are opt-in in CuTe +// and must be enabled before the first CuTe include of the translation unit. +// Including this header first is what guarantees that. +#if !defined(CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED) +#define CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +namespace liger { +namespace sm90 { + +using namespace cute; + +// ─────────────────────────────────────────────────────────────────────────── +// Architecture constants +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ArchSm90 { + static_assert(Compute == 90, "ArchSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + static constexpr int kWarpSize = 32; + static constexpr int kWarpsPerWarpGroup = 4; + static constexpr int kWarpGroupSize = kWarpSize * kWarpsPerWarpGroup; + // Hopper's opt-in dynamic shared-memory ceiling. + static constexpr int kMaxDynamicSmemBytes = 227 * 1024; + // A CTA may own at most 64K registers; ptxas therefore caps a 384-thread + // warp-specialized CTA at 168 registers per thread. + static constexpr int kRegisterFilePerCta = 65536; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// WGMMA atom / TiledMMA selection +// ─────────────────────────────────────────────────────────────────────────── + +// The M64 SS WGMMA atom for the requested N. Wraps CuTe's selector so callers +// do not have to know which shapes live in the extended set. +template < + class ElementA, + class ElementB, + class ElementAccum, + int AtomTileN, + GMMA::Major MajorA = GMMA::Major::K, + GMMA::Major MajorB = GMMA::Major::K, + int Compute = 90> +struct GmmaAtomSm90 { + static_assert(Compute == 90, "GmmaAtomSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + static constexpr int kAtomTileM = 64; + static constexpr int kAtomTileN = AtomTileN; + static constexpr int kAtomTileK = 16; + + using Atom = decltype(cute::GMMA::ss_op_selector< + ElementA, + ElementB, + ElementAccum, + Shape, Int, Int>, + MajorA, + MajorB>()); +}; + +// Cooperative TiledMMA: NumWarpGroups warp groups split the tile along M, so +// each warp group owns an (AtomTileM, AtomTileN) accumulator. +template < + class ElementA, + class ElementB, + class ElementAccum, + int AtomTileN, + int NumWarpGroups, + GMMA::Major MajorA = GMMA::Major::K, + GMMA::Major MajorB = GMMA::Major::K, + int Compute = 90> +struct TiledMmaMSplitSm90 { + static_assert(Compute == 90, "TiledMmaMSplitSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + static constexpr int kNumWarpGroups = NumWarpGroups; + + using AtomTraits = GmmaAtomSm90< + ElementA, ElementB, ElementAccum, AtomTileN, MajorA, MajorB, Compute>; + using Atom = typename AtomTraits::Atom; + using Type = TiledMMA< + MMA_Atom, Layout, _1, _1>>>; + + static constexpr int kTileM = AtomTraits::kAtomTileM * NumWarpGroups; + static constexpr int kTileN = AtomTileN; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// SMEM layouts +// ─────────────────────────────────────────────────────────────────────────── + +// K-major operand staging with the 128 B swizzle atom, which is what a K tile +// of exactly 128 bytes requires. +template +struct SmemLayoutKMajorSm90 { + static_assert(Compute == 90, "SmemLayoutKMajorSm90 requires Compute=90"); + static_assert(Cols * static_cast(sizeof(Element)) % 128 == 0, + "the SW128 K-major atom needs 128 B rows"); + + static constexpr int kCompute = Compute; + + using Atom = GMMA::Layout_K_SW128_Atom; + using Single = decltype(tile_to_shape( + Atom{}, Shape, Int>{})); + using Staged = decltype(tile_to_shape( + Atom{}, Shape, Int, Int>{})); + + static constexpr int kStageBytes = static_cast( + cute::size(Single{}) * sizeof(Element)); + static constexpr int kStagedElements = cosize_v; +}; + +// ─────────────────────────────────────────────────────────────────────────── +// TMA descriptors +// ─────────────────────────────────────────────────────────────────────────── + +template +struct TmaSm90 { + static_assert(Compute == 90, "TmaSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + + // int64 extents: rows * cols exceeds INT_MAX at production shapes and + // would otherwise wrap CuTe's layout algebra into a negative byte offset. + template + static auto row_major_tensor(const Element* ptr, int rows, int cols) { + return make_tensor( + make_gmem_ptr(ptr), + make_layout( + make_shape( + static_cast(rows), static_cast(cols)), + make_stride(static_cast(cols), _1{}))); + } + + template + static auto make_load(const Element* ptr, int rows, int cols) { + return make_tma_copy( + SM90_TMA_LOAD{}, + row_major_tensor(ptr, rows, cols), + SmemLayout{}, + make_shape(Int{}, Int{}), + _1{}); + } + + // Multicast splits the box across `Multicast` CTAs of the cluster; each + // peer issues its own slice and every peer receives the whole box, so the + // stage's transaction byte count stays the full tile. + template < + class SmemLayout, int TileRows, int TileCols, int Multicast, + class Element> + static auto make_load_multicast(const Element* ptr, int rows, int cols) { + return make_tma_copy( + SM90_TMA_LOAD_MULTICAST{}, + row_major_tensor(ptr, rows, cols), + SmemLayout{}, + make_shape(Int{}, Int{}), + Int{}); + } +}; + +// Multicast mask over the cluster's M mode: every CTA sharing this N tile. +template +CUTE_DEVICE std::uint16_t cluster_multicast_mask_m_sm90() { + static_assert(Compute == 90, + "cluster_multicast_mask_m_sm90 requires Compute=90"); + std::uint16_t mask = 0; + auto block_layout = Layout{}; + CUTE_UNROLL + for (int m = 0; m < size<0>(block_layout); ++m) { + mask |= std::uint16_t(1) << block_layout(m, Int<0>{}, Int<0>{}); + } + return mask; +} + +// ─────────────────────────────────────────────────────────────────────────── +// Mainloop pipeline +// ─────────────────────────────────────────────────────────────────────────── + +enum class PipelineRoleSm90 : std::uint8_t { + kProducer, + kConsumer, + kNonParticipant, +}; + +template +struct MainloopPipelineSm90 { + static_assert(Compute == 90, "MainloopPipelineSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + static constexpr int kStages = Stages; + + using Type = cutlass::PipelineTmaAsync; + using State = cutlass::PipelineState; + using SharedStorage = typename Type::SharedStorage; + using Category = typename Type::ThreadCategory; + + // Role-specific construction. Barrier initialisation is split out of the + // constructors so a warp-specialized kernel can (a) reconfigure registers + // before any pipeline object exists and (b) keep the producer's and the + // consumer's pipeline objects in disjoint scopes — neither is ever live on + // the other's path, and idle warps never materialise one at all. + + // All threads call this; only `initializing_warp` touches the mbarriers, + // but every thread must issue the init fence. Must be followed by a + // cluster-wide handshake before any barrier is used. + template + CUTE_DEVICE static void init_barriers( + SharedStorage& storage, + int transaction_bytes, + int num_consumers, + ClusterShape cluster_shape) { + typename Type::Params params; + params.transaction_bytes = transaction_bytes; + params.num_producers = 1; + params.num_consumers = num_consumers; + Type::init_barriers(storage, params, cluster_shape); + } + + template + CUTE_DEVICE static Type make_producer( + SharedStorage& storage, + int transaction_bytes, + int num_consumers, + bool is_leader, + ClusterShape cluster_shape) { + typename Type::Params params; + params.transaction_bytes = transaction_bytes; + params.num_producers = 1; + params.num_consumers = num_consumers; + params.is_leader = is_leader; + params.role = Category::Producer; + // Barriers are already live; only the per-thread masks are built here. + return Type(storage, params, cluster_shape, + cute::false_type{}, cute::true_type{}); + } + + template + CUTE_DEVICE static Type make_consumer( + SharedStorage& storage, + int transaction_bytes, + int num_consumers, + ClusterShape cluster_shape) { + typename Type::Params params; + params.transaction_bytes = transaction_bytes; + params.num_producers = 1; + params.num_consumers = num_consumers; + params.is_leader = 0; + params.role = Category::Consumer; + return Type(storage, params, cluster_shape, + cute::false_type{}, cute::true_type{}); + } + + CUTE_DEVICE static State producer_start_state() { + return cutlass::make_producer_start_state(); + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Warp-group register reconfiguration and cluster synchronisation +// ─────────────────────────────────────────────────────────────────────────── + +template +struct WarpSpecializedRegistersSm90 { + static_assert(Compute == 90, + "WarpSpecializedRegistersSm90 requires Compute=90"); + static_assert(ProducerRegisters % 8 == 0 && ConsumerRegisters % 8 == 0, + "setmaxnreg operands must be multiples of 8"); + + static constexpr int kCompute = Compute; + static constexpr int kProducerRegisters = ProducerRegisters; + static constexpr int kConsumerRegisters = ConsumerRegisters; + + // Must be executed by every warp of the warp group, including idle ones. + CUTE_DEVICE static void producer() { + cutlass::arch::warpgroup_reg_dealloc(); + } + + CUTE_DEVICE static void consumer() { + cutlass::arch::warpgroup_reg_alloc(); + } +}; + +// pipeline_init_arrive / pipeline_init_wait: the cluster-wide handshake that +// must complete before any mbarrier in the cluster is touched. +template +CUTE_DEVICE void cluster_pipeline_init_sm90() { + static_assert(Compute == 90, + "cluster_pipeline_init_sm90 requires Compute=90"); + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); +} + +// No CTA may retire while a cluster peer can still address its SMEM. +template +CUTE_DEVICE void cluster_exit_sm90() { + static_assert(Compute == 90, "cluster_exit_sm90 requires Compute=90"); + cute::cluster_sync(); +} + +template +CUTE_DEVICE void named_barrier_sync_sm90() { + static_assert(Compute == 90, "named_barrier_sync_sm90 requires Compute=90"); + static_assert(BarrierId >= 0 && BarrierId < 16, + "hardware named barrier ids are 0..15"); + static_assert(NumThreads % 32 == 0, + "named barrier thread counts must be warp aligned"); + asm volatile("barrier.sync %0, %1;" + : + : "n"(BarrierId), "n"(NumThreads)); +} + +// ─────────────────────────────────────────────────────────────────────────── +// Mainloop drivers +// +// These encode the CuTe-DSL pipeline *ordering* and nothing else. The caller +// supplies what to issue per stage; both functors are inlined, so the emitted +// instruction stream is identical to writing the loop out by hand. +// ─────────────────────────────────────────────────────────────────────────── + +template +struct MainloopKLoopSm90 { + static_assert(Compute == 90, "MainloopKLoopSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + + // Producer: acquire the stage, issue every TMA copy for it, advance. + // `issue(barrier, stage, k_tile)` performs the copies. + template + CUTE_DEVICE static void produce( + Pipeline& pipe, + State& state, + int num_k_tiles, + IssueCopies issue) { + for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) { + pipe.producer_acquire(state); + issue(pipe.producer_get_barrier(state), state.index(), k_tile); + ++state; + } + } + + // Consumer with the deliberate one-stage release lag. + // + // Stage zero is issued without a release because no earlier WGMMA group + // has retired. Every later stage is committed, `wait_group(1)` retires the + // previous group, and only then is that previous AB stage handed back to + // the producer. The tail uses `wait_group(0)` before the final release. + // Without the lag, TMA would overwrite operands still feeding a live + // WGMMA group. + // + // `issue_mma(stage)` fences the accumulators, arrives, issues the GEMMs + // and commits the batch. `fence_accumulators()` re-fences them after a + // wait_group so the accumulator registers are readable again. + template + CUTE_DEVICE static void consume( + Pipeline& pipe, + State& read_state, + State& release_state, + int num_k_tiles, + IssueMma issue_mma, + FenceAccumulators fence_accumulators) { + pipe.consumer_wait(read_state); + issue_mma(read_state.index()); + ++read_state; + + for (int k_tile = 1; k_tile < num_k_tiles; ++k_tile) { + pipe.consumer_wait(read_state); + issue_mma(read_state.index()); + + warpgroup_wait<1>(); + fence_accumulators(); + pipe.consumer_release(release_state); + ++release_state; + ++read_state; + } + + warpgroup_wait<0>(); + fence_accumulators(); + pipe.consumer_release(release_state); + ++release_state; + } +}; + +// ─────────────────────────────────────────────────────────────────────────── +// Host-side cluster launch +// ─────────────────────────────────────────────────────────────────────────── + +template +struct ClusterLaunchSm90 { + static_assert(Compute == 90, "ClusterLaunchSm90 requires Compute=90"); + + static constexpr int kCompute = Compute; + + // Opt into the Hopper dynamic-SMEM ceiling and non-portable cluster sizes. + template + static cudaError_t prepare(Kernel kernel, int smem_bytes) { + cudaError_t error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + if (error != cudaSuccess) return error; + return cudaFuncSetAttribute( + kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); + } + + static cudaLaunchConfig_t make_config( + dim3 grid, + int threads, + int smem_bytes, + int cluster_m, + cudaLaunchAttribute& attribute, + cudaStream_t stream) { + attribute = {}; + attribute.id = cudaLaunchAttributeClusterDimension; + attribute.val.clusterDim.x = static_cast(cluster_m); + attribute.val.clusterDim.y = 1; + attribute.val.clusterDim.z = 1; + + cudaLaunchConfig_t config = {}; + config.gridDim = grid; + config.blockDim = dim3(static_cast(threads), 1, 1); + config.dynamicSmemBytes = smem_bytes; + config.stream = stream; + config.attrs = &attribute; + config.numAttrs = 1; + return config; + } + + // Concurrent clusters this device can hold, i.e. the occupancy target a + // persistent/split scheduler should aim at. Returns 0 if unavailable. + template + static int max_active_clusters( + Kernel kernel, int threads, int smem_bytes, int cluster_m) { + cudaLaunchAttribute attribute = {}; + cudaLaunchConfig_t config = make_config( + dim3(static_cast(cluster_m), 1, 1), + threads, + smem_bytes, + cluster_m, + attribute, + nullptr); + int clusters = 0; + if (cudaOccupancyMaxActiveClusters(&clusters, kernel, &config) != + cudaSuccess) { + cudaGetLastError(); + return 0; + } + return clusters; + } + + template + static cudaError_t launch( + Kernel kernel, + dim3 grid, + int threads, + int smem_bytes, + int cluster_m, + cudaStream_t stream, + Args const&... args) { + cudaLaunchAttribute attribute = {}; + cudaLaunchConfig_t config = make_config( + grid, threads, smem_bytes, cluster_m, attribute, stream); + return cudaLaunchKernelEx(&config, kernel, args...); + } +}; + +} // namespace sm90 +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/online_softmax.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/online_softmax.cuh new file mode 100644 index 00000000..666c9a7f --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/online_softmax.cuh @@ -0,0 +1,118 @@ +#pragma once + +#include + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +struct OnlineSoftmaxState { + float max_value; + float exp_sum; + float exp_weighted_sum; + float target_logit; + int has_target; + + __host__ __device__ OnlineSoftmaxState() + : max_value(-1.0e38f), + exp_sum(0.0f), + exp_weighted_sum(0.0f), + target_logit(0.0f), + has_target(0) {} + + template + __host__ __device__ void fold(float logit, bool valid, bool is_target) { + if (!valid) return; + + float next_max = fmaxf(max_value, logit); + float previous_scale = exp_sum == 0.0f + ? 0.0f + : expf(max_value - next_max); + float logit_scale = expf(logit - next_max); + + exp_sum = exp_sum * previous_scale + logit_scale; + if constexpr (ReturnEntropy) { + exp_weighted_sum = + exp_weighted_sum * previous_scale + logit_scale * logit; + } + max_value = next_max; + + if (is_target) { + target_logit = logit; + has_target = 1; + } + } + +}; + +template +__host__ __device__ inline OnlineSoftmaxState merge_online_softmax( + const OnlineSoftmaxState& lhs, const OnlineSoftmaxState& rhs) { + if (lhs.exp_sum == 0.0f) return rhs; + if (rhs.exp_sum == 0.0f) return lhs; + + OnlineSoftmaxState result; + result.max_value = fmaxf(lhs.max_value, rhs.max_value); + float lhs_scale = expf(lhs.max_value - result.max_value); + float rhs_scale = expf(rhs.max_value - result.max_value); + result.exp_sum = lhs.exp_sum * lhs_scale + rhs.exp_sum * rhs_scale; + if constexpr (ReturnEntropy) { + result.exp_weighted_sum = + lhs.exp_weighted_sum * lhs_scale + rhs.exp_weighted_sum * rhs_scale; + } + result.has_target = lhs.has_target || rhs.has_target; + result.target_logit = lhs.has_target ? lhs.target_logit : rhs.target_logit; + return result; +} + +struct CorrectedSoftmaxStats { + float exp_sum; + float exp_weighted_sum; + float target_logit; +}; + +template +__host__ __device__ inline CorrectedSoftmaxStats correct_for_global_max( + const OnlineSoftmaxState& local, float global_max) { + if (local.exp_sum == 0.0f) { + return {0.0f, 0.0f, local.has_target ? local.target_logit : 0.0f}; + } + float scale = expf(local.max_value - global_max); + float corrected_weighted_sum = 0.0f; + if constexpr (ReturnEntropy) { + corrected_weighted_sum = local.exp_weighted_sum * scale; + } + return { + local.exp_sum * scale, + corrected_weighted_sum, + local.has_target ? local.target_logit : 0.0f, + }; +} + +struct FinalizedSoftmax { + float lse; + float nll; + float entropy; +}; + +template +__host__ __device__ inline FinalizedSoftmax finalize_softmax( + float global_max, + float global_sum, + float global_target_logit, + float global_weighted_sum, + bool ignored) { + float lse = global_max + logf(global_sum); + float nll = ignored ? 0.0f : lse - global_target_logit; + float entropy = 0.0f; + if constexpr (ReturnEntropy) { + if (!ignored) { + entropy = lse - global_weighted_sum / global_sum; + } + } + return {lse, nll, entropy}; +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/state.h b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/state.h new file mode 100644 index 00000000..5b81385c --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/state.h @@ -0,0 +1,11 @@ +#pragma once + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +// Invalidates cached FSLCE capacities, mappings, and duplicated teams before +// the shared BufferPool is cleared or NVSHMEM is finalized. +void reset_fslce_tp_configuration(); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/tma_store_sm90.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/tma_store_sm90.cuh new file mode 100644 index 00000000..ccd4a948 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/tma_store_sm90.cuh @@ -0,0 +1,36 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// SM90 fused scaled linear cross entropy — TMA store completion support. +// +// Transport selection and collective mappings live in liger_cute/detail. +// ═══════════════════════════════════════════════════════════════════════════ + +#include + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +// ─────────────────────────────────────────────────────────────────────────── +// TMA store completion helpers used by the dZ and dW epilogues. +// ─────────────────────────────────────────────────────────────────────────── + +// cp.async.bulk.wait_group N, always the full form and never `.read`. +// +// `.read` only waits for the SMEM source of a bulk store to be re-usable, and +// it *retires the group from the pending set*. A staging pipeline that drains +// with `.read` and then asks for `wait_group 0` therefore gets no guarantee at +// all about the already-retired groups' global writes. Every tile here is +// published to another warp the moment it is stored, so the epilogue has to +// wait for completion, not merely for the source to be free. +template +__device__ __forceinline__ void wait_tma_store() { + asm volatile("cp.async.bulk.wait_group %0;" :: "n"(Pending) : "memory"); +} + +__device__ __forceinline__ void wait_tma_store_complete() { + wait_tma_store<0>(); +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cu b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cu new file mode 100644 index 00000000..84069082 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cu @@ -0,0 +1,429 @@ +#include "workspace.cuh" + +#include + +#include +#include +#include + +#include "buffer_pool.cuh" +#include "forward_reduce.cuh" +#include "liger_cute/check.h" +#include "liger_cute/detail/tp_reduce.cuh" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { +namespace { + +using Config = BackwardGemmConfigSm90<90>; +using Launch = BackwardGemmLaunchSm90<90>; + +// The ring depth is a compile-time constant shared with the kernel template +// (dx_reduce.cuh): the symmetric capacity and CTA-local mbarrier indexing have +// to match what the kernel uses. +constexpr int kConfiguredStages = kDxRingStages; + +BackwardTpCapacity g_capacity = {}; +bool g_configured = false; +std::size_t g_staging_bytes = 0; +std::size_t g_durable_bytes = 0; +std::size_t g_packed_durable_bytes = 0; +std::size_t g_reduced_bytes = 0; + +std::size_t checked_multiply(std::size_t lhs, std::size_t rhs, const char* name) { + LIGER_CHECK( + rhs == 0 || lhs <= std::numeric_limits::max() / rhs, + name, + " size overflow"); + return lhs * rhs; +} + +std::size_t staging_bytes_at(int tiles_per_reduce, int num_ctas) { + std::size_t tile_elements = checked_multiply( + static_cast(Config::kTileM), + static_cast(Config::kDxTileN), + "dX tile"); + std::size_t elements = checked_multiply( + tile_elements, + static_cast(tiles_per_reduce), + "dX group"); + elements = checked_multiply( + elements, static_cast(kConfiguredStages), "dX ring"); + elements = checked_multiply( + elements, + static_cast(kDxCommWarpsPerChannel), + "dX comm warps"); + elements = checked_multiply( + elements, static_cast(num_ctas), "dX resident CTAs"); + return checked_multiply(elements, sizeof(float), "dX staging"); +} + +std::size_t sync_bytes_at(int num_ctas, int team_size) { + std::size_t entries = checked_multiply( + static_cast(num_ctas), + static_cast(kDxCommWarpsPerChannel), + "dX sync CTAs"); + entries = checked_multiply( + entries, static_cast(kConfiguredStages), "dX sync ring"); + entries = checked_multiply( + entries, static_cast(kDxSyncPhases), "dX sync phases"); + entries = checked_multiply( + entries, static_cast(team_size), "dX sync team"); + return checked_multiply(entries, sizeof(std::uint64_t), "dX sync"); +} + +std::size_t durable_bytes_at(int tokens, int hidden) { + LIGER_CHECK(tokens > 0, "max_tokens must be positive"); + LIGER_CHECK(hidden > 0, "max_hidden must be positive"); + std::size_t waves = + (static_cast(tokens) + Config::kWaveRows - 1) / + Config::kWaveRows; + std::size_t rows = checked_multiply( + waves, static_cast(Config::kWaveRows), "dX durable rows"); + std::size_t n_tiles = + (static_cast(hidden) + Config::kDxTileN - 1) / + Config::kDxTileN; + std::size_t columns = checked_multiply( + n_tiles, + static_cast(Config::kDxTileN), + "dX durable columns"); + return checked_multiply( + checked_multiply(rows, columns, "dX durable elements"), + sizeof(float), + "dX durable"); +} + +void ensure_configured() { + LIGER_CHECK( + g_configured, + "fused_scaled_linear_cross_entropy backward: call " + "configure_backward_tp_symmetric(max_tokens, max_hidden, " + "max_local_vocab, " + "max_tiles_per_reduce, max_comm_channels, team_handle) collectively on " + "every PE before the first launch"); +} + +void check_cuda(cudaError_t error, const char* what) { + LIGER_CHECK( + error == cudaSuccess, + "fused_scaled_linear_cross_entropy backward: ", + what, + " failed: ", + cudaGetErrorString(error)); +} + +int resident_cta_capacity() { + int device = 0; + check_cuda(cudaGetDevice(&device), "cudaGetDevice"); + cudaDeviceProp properties = {}; + check_cuda( + cudaGetDeviceProperties(&properties, device), + "cudaGetDeviceProperties"); + LIGER_CHECK( + properties.multiProcessorCount > 0 && + properties.multiProcessorCount <= kMaxDxResidentCtas, + "fused_scaled_linear_cross_entropy backward: SM90 device reports ", + properties.multiProcessorCount, + " SMs, but the CTA-owned dX workspace supports at most ", + kMaxDxResidentCtas); + return properties.multiProcessorCount; +} + +} // namespace + +void configure_backward_tp_symmetric( + int max_tokens, + int max_hidden, + int max_local_vocab, + int max_tiles_per_reduce, + int max_comm_channels, + std::int64_t team_handle) { + LIGER_CHECK(max_tokens > 0, "max_tokens must be positive"); + LIGER_CHECK(max_hidden > 0, "max_hidden must be positive"); + LIGER_CHECK(max_local_vocab > 0, "max_local_vocab must be positive"); + LIGER_CHECK( + max_tiles_per_reduce >= 1, + "max_tiles_per_reduce must be positive"); + LIGER_CHECK( + max_comm_channels >= 1, + "max_comm_channels must be positive"); + + int max_resident_ctas = resident_cta_capacity(); + liger_cute::detail::TpReduceTopology topology = + liger_cute::detail::query_tp_reduce_topology(team_handle); + int team_size = topology.team_size; + + if (g_configured) { + LIGER_CHECK( + max_tokens <= g_capacity.max_tokens && + max_hidden <= g_capacity.max_hidden && + max_local_vocab <= g_capacity.max_local_vocab && + max_tiles_per_reduce <= g_capacity.max_tiles_per_reduce && + max_comm_channels <= g_capacity.max_comm_channels && + max_resident_ctas == g_capacity.max_resident_ctas && + team_size == g_capacity.team_size && + team_handle == g_capacity.team_handle, + "fused_scaled_linear_cross_entropy backward: the symmetric " + "capacity is immutable once allocated (configured for vocab ", + g_capacity.max_local_vocab, + ", tokens ", + g_capacity.max_tokens, + ", hidden ", + g_capacity.max_hidden, + ", TilesPerReduce ", + g_capacity.max_tiles_per_reduce, + ", communication channels ", + g_capacity.max_comm_channels, + ", resident CTAs ", + g_capacity.max_resident_ctas, + "). Configure the maximum upfront."); + return; + } + + using Names = BackwardSymmetricNames; + auto& pool = global_buffer_pool(); + + // One fixed order on every PE: symmetric first, then device private. + g_staging_bytes = staging_bytes_at( + max_tiles_per_reduce, max_resident_ctas); + g_durable_bytes = durable_bytes_at(max_tokens, max_hidden); + int local_partition_size = topology.local_size; + g_packed_durable_bytes = + (g_durable_bytes + static_cast(local_partition_size) - 1) / + static_cast(local_partition_size); + g_reduced_bytes = local_partition_size > 1 + ? (g_staging_bytes > g_durable_bytes + ? g_staging_bytes + : g_durable_bytes) + : g_staging_bytes; + auto* partial = static_cast( + pool.get_symmetric(Names::kDxPartial, g_staging_bytes)); + auto* reduced = static_cast( + pool.get_symmetric(Names::kDxReduced, g_reduced_bytes)); + auto* reduced_shard = static_cast( + pool.get_symmetric( + Names::kDxReducedShard, g_packed_durable_bytes)); + auto* remote_inbox = static_cast( + pool.get_symmetric( + Names::kDxRemoteInbox, g_packed_durable_bytes)); + auto* remote_signals = static_cast( + pool.get_symmetric( + Names::kDxRemoteSignals, + 2 * sizeof(std::uint64_t))); + std::size_t sync_bytes = sync_bytes_at(max_resident_ctas, team_size); + auto* sync = static_cast( + pool.get_symmetric(Names::kDxSync, sync_bytes)); + auto** peer_partial_storage = static_cast( + pool.get_device( + Names::kDxPeerPartialPointers, + static_cast(team_size) * sizeof(float*))); + auto** peer_sync_storage = static_cast( + pool.get_device( + Names::kDxPeerSyncPointers, + static_cast(team_size) * + sizeof(std::uint64_t*))); + pool.get_device( + Names::kDzWorkspace, Launch::dz_workspace_bytes(max_local_vocab)); + auto* launch_epoch = static_cast( + pool.get_device(Names::kDxLaunchEpoch, sizeof(std::uint64_t))); + check_cuda( + cudaMemset(launch_epoch, 0, sizeof(std::uint64_t)), + "cudaMemset(dX launch epoch)"); + + liger_cute::detail::configure_tp_reduce( + team_handle, + { + partial, + reduced, + reduced_shard, + remote_inbox, + remote_signals, + sync, + sync_bytes, + peer_partial_storage, + peer_sync_storage, + }); + g_capacity.max_tokens = max_tokens; + g_capacity.max_hidden = max_hidden; + g_capacity.max_local_vocab = max_local_vocab; + g_capacity.max_tiles_per_reduce = max_tiles_per_reduce; + g_capacity.max_comm_channels = max_comm_channels; + g_capacity.max_resident_ctas = max_resident_ctas; + g_capacity.max_stages = kConfiguredStages; + g_capacity.team_size = team_size; + g_capacity.team_handle = team_handle; + g_configured = true; +} + +std::size_t backward_tp_pool_symmetric_bytes( + int max_tokens, + int max_hidden, + int max_tiles_per_reduce, + int max_comm_channels) { + LIGER_CHECK( + max_comm_channels >= 1, + "max_comm_channels must be positive"); + int max_ctas = + g_configured ? g_capacity.max_resident_ctas : resident_cta_capacity(); + int team_size = g_configured + ? g_capacity.team_size + : liger_cute::detail::kMaxTpReduceTeamSize; + if (g_configured) { + return g_staging_bytes + g_reduced_bytes + + 2u * g_packed_durable_bytes + + sync_bytes_at(max_ctas, team_size) + + 2 * sizeof(std::uint64_t); + } + std::size_t staging = staging_bytes_at(max_tiles_per_reduce, max_ctas); + std::size_t durable = durable_bytes_at(max_tokens, max_hidden); + std::size_t reduced = staging > durable ? staging : durable; + return staging + reduced + 2u * durable + + sync_bytes_at(max_ctas, team_size) + + 2 * sizeof(std::uint64_t); +} + +std::size_t backward_tp_pool_device_bytes(int max_local_vocab) { + int team_size = g_configured + ? g_capacity.team_size + : liger_cute::detail::kMaxTpReduceTeamSize; + std::size_t peer_mapping = + static_cast(team_size) * + (2 * sizeof(void*)); + return Launch::dz_workspace_bytes(max_local_vocab) + + sizeof(std::uint64_t) + peer_mapping; +} + +DxReduceWorkspace reserve_dx_reduce_workspace( + int tiles_per_reduce, + int num_stages, + int num_ctas) { + ensure_configured(); + LIGER_CHECK( + tiles_per_reduce >= 1 && + tiles_per_reduce <= g_capacity.max_tiles_per_reduce, + "fused_scaled_linear_cross_entropy backward: TilesPerReduce ", + tiles_per_reduce, + " exceeds the configured maximum ", + g_capacity.max_tiles_per_reduce); + LIGER_CHECK( + num_stages == g_capacity.max_stages, + "fused_scaled_linear_cross_entropy backward: the staging ring depth ", + num_stages, + " must equal the configured ", + g_capacity.max_stages, + "; CTA-local mbarriers use the compile-time depth"); + LIGER_CHECK( + num_ctas >= 1 && num_ctas <= g_capacity.max_resident_ctas, + "fused_scaled_linear_cross_entropy backward: launch grid ", + num_ctas, + " exceeds resident CTA capacity ", + g_capacity.max_resident_ctas); + using Names = BackwardSymmetricNames; + auto& pool = global_buffer_pool(); + // Always the configured capacity, never the per-call size. + std::size_t sync_bytes = sync_bytes_at( + g_capacity.max_resident_ctas, g_capacity.team_size); + DxReduceWorkspace workspace = {}; + workspace.partial = static_cast( + pool.get_symmetric(Names::kDxPartial, g_staging_bytes)); + workspace.reduced = static_cast( + pool.get_symmetric(Names::kDxReduced, g_reduced_bytes)); + workspace.sync = static_cast( + pool.get_symmetric(Names::kDxSync, sync_bytes)); + workspace.launch_epoch = static_cast( + pool.get_device(Names::kDxLaunchEpoch, sizeof(std::uint64_t))); + return workspace; +} + +std::size_t backward_dx_staging_bytes(int max_tiles_per_reduce) { + int max_ctas = + g_configured ? g_capacity.max_resident_ctas : resident_cta_capacity(); + return staging_bytes_at(max_tiles_per_reduce, max_ctas); +} + +std::size_t backward_dx_configured_staging_bytes() { + ensure_configured(); + return staging_bytes_at( + g_capacity.max_tiles_per_reduce, g_capacity.max_resident_ctas); +} + +std::size_t backward_dx_configured_durable_bytes() { + ensure_configured(); + return g_durable_bytes; +} + +void validate_backward_tp_shape( + int tokens, int hidden, int local_vocab) { + ensure_configured(); + LIGER_CHECK( + tokens > 0 && tokens <= g_capacity.max_tokens, + "fused_scaled_linear_cross_entropy backward: tokens ", + tokens, + " exceed the configured maximum ", + g_capacity.max_tokens); + LIGER_CHECK( + hidden > 0 && hidden <= g_capacity.max_hidden, + "fused_scaled_linear_cross_entropy backward: hidden ", + hidden, + " exceeds the configured maximum ", + g_capacity.max_hidden); + LIGER_CHECK( + local_vocab > 0 && local_vocab <= g_capacity.max_local_vocab, + "fused_scaled_linear_cross_entropy backward: local_vocab ", + local_vocab, + " exceeds the configured maximum ", + g_capacity.max_local_vocab); + LIGER_CHECK( + durable_bytes_at(tokens, hidden) <= + backward_dx_configured_durable_bytes(), + "fused_scaled_linear_cross_entropy backward: cluster dX durable " + "workspace capacity exceeded"); +} + +int backward_dx_resident_cta_capacity() { + return g_configured ? g_capacity.max_resident_ctas : resident_cta_capacity(); +} + +int backward_dx_team_size() { + ensure_configured(); + return g_capacity.team_size; +} + +std::int64_t backward_dx_team_handle() { + ensure_configured(); + return g_capacity.team_handle; +} + +void reset_fslce_tp_configuration() { + liger_cute::detail::reset_tp_reduce(); + g_capacity = {}; + g_configured = false; + g_staging_bytes = 0; + g_durable_bytes = 0; + g_packed_durable_bytes = 0; + g_reduced_bytes = 0; + reset_forward_tp_workspace_configuration(); +} + +BackwardScratch reserve_backward_scratch(int local_vocab) { + ensure_configured(); + LIGER_CHECK( + local_vocab > 0 && local_vocab <= g_capacity.max_local_vocab, + "fused_scaled_linear_cross_entropy backward: local_vocab ", + local_vocab, + " exceeds the configured maximum ", + g_capacity.max_local_vocab); + + using Names = BackwardSymmetricNames; + auto& pool = global_buffer_pool(); + std::size_t dz_bytes = + Launch::dz_workspace_bytes(g_capacity.max_local_vocab); + + BackwardScratch scratch = {}; + scratch.dz_workspace = pool.get_device(Names::kDzWorkspace, dz_bytes); + scratch.dz_workspace_bytes = dz_bytes; + return scratch; +} + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cuh b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cuh new file mode 100644 index 00000000..37408f31 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cuh @@ -0,0 +1,147 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════════════ +// Backward workspace ownership. +// +// ── Allocator ownership ─────────────────────────────────────────────────── +// Nothing here allocates memory itself: there is no local allocator and no raw +// nvshmem_malloc / nvshmem_free lifecycle anywhere in the backward path. Every +// buffer comes from liger::global_buffer_pool(), which is the shared MoE +// utility for fixed, reusable, internal staging that never escapes the +// operation. +// +// name allocator symmetric? why +// ----------------- --------- ---------- --------------------------- +// kDxPartial pool yes FP32 NVLS SUM source +// kDxReduced pool yes FP32 NVLS SUM destination +// kDxSync pool yes per-slot ready/completion epochs +// kDzWorkspace pool no one wave of dZ, BF16 +// +// liger::global_symmetric_stack() is deliberately NOT used. Nothing produced by +// the forward has to be carried to the backward as symmetric state: the +// forward's two all-reduces leave every rank holding identical NLL / LSE / +// entropy, so those are ordinary local CUDA outputs saved through autograd. +// Backward consumes them exactly like any other saved tensor. +// +// ── Collective configuration ────────────────────────────────────────────── +// configure_backward_tp_symmetric() must run on every PE with the same values +// before the first launch. It fixes the capacity of every allocation named +// here and resolves their aliases in the TP team's multicast mapping. +// Capacities are immutable afterwards: every later request asks for the +// *configured* capacity, never the per-call size, so the byte count handed to +// the pool is identical on every call and on every PE. +// ═══════════════════════════════════════════════════════════════════════════ + +#include +#include + +#include +#include + +#include "backward_gemm_sm90.cuh" +#include "dx_reduce.cuh" +#include "state.h" + +namespace liger { +namespace fused_scaled_linear_cross_entropy { + +// Centralized allocation names. Nothing else in the backward path may name a +// pooled buffer. +struct BackwardSymmetricNames { + // global_buffer_pool() — symmetric (remotely addressed by NVSHMEM). + static constexpr const char* kDxPartial = + "fused_scaled_linear_cross_entropy_tp_dx_partial"; + static constexpr const char* kDxReduced = + "fused_scaled_linear_cross_entropy_tp_dx_reduced"; + static constexpr const char* kDxReducedShard = + "fused_scaled_linear_cross_entropy_tp_dx_reduced_shard"; + static constexpr const char* kDxRemoteInbox = + "fused_scaled_linear_cross_entropy_tp_dx_remote_inbox"; + static constexpr const char* kDxRemoteSignals = + "fused_scaled_linear_cross_entropy_tp_dx_remote_signals"; + static constexpr const char* kDxSync = + "fused_scaled_linear_cross_entropy_tp_dx_sync"; + + // global_buffer_pool() — device private (never remotely addressed). + static constexpr const char* kDzWorkspace = + "fused_scaled_linear_cross_entropy_tp_dz_workspace"; + static constexpr const char* kDxPeerPartialPointers = + "fused_scaled_linear_cross_entropy_tp_dx_peer_partial_ptrs"; + static constexpr const char* kDxPeerSyncPointers = + "fused_scaled_linear_cross_entropy_tp_dx_peer_sync_ptrs"; + static constexpr const char* kDxLaunchEpoch = + "fused_scaled_linear_cross_entropy_tp_dx_launch_epoch"; +}; + +// The immutable collective configuration. +struct BackwardTpCapacity { + int max_tokens; + int max_hidden; + int max_local_vocab; + int max_tiles_per_reduce; + int max_comm_channels; + int max_resident_ctas; + int max_stages; + int team_size; + std::int64_t team_handle; +}; + +// Collective. Must be called on every PE with the same values before the first +// launch. CTA-owned NVLS staging is sized for full residency; +// `max_comm_channels` is retained for API compatibility; the CTA-owned +// production path does not allocate per-channel teams or buffers. +void configure_backward_tp_symmetric( + int max_tokens, + int max_hidden, + int max_local_vocab, + int max_tiles_per_reduce, + int max_comm_channels, + std::int64_t team_handle); + +// Symmetric footprint at the configured maximum. The fused path reuses the +// larger CTA-owned staging allocation, so channels add no buffer bytes. +std::size_t backward_tp_pool_symmetric_bytes( + int max_tokens, + int max_hidden, + int max_tiles_per_reduce, + int max_comm_channels); + +// Total device-private pool footprint. CTA ring state is in shared memory. +std::size_t backward_tp_pool_device_bytes(int max_local_vocab); + +// Collective: every PE of the configured team must call this with the same +// arguments. Returns local and multicast staging pointers plus TP metadata. +// Every pool request asks for the *configured* capacity, +// so the byte count is identical on every PE and on every call; only the +// descriptive fields of the returned workspace follow the launch grid. +// +// The knobs are runtime here on purpose: the compile-time TilesPerReduce / +// NumStages of the kernel template are validated against them in the launcher, +// which keeps a single allocation path for every instantiation. +DxReduceWorkspace reserve_dx_reduce_workspace( + int tiles_per_reduce, int num_stages, int num_ctas); + +// Bytes in one of the two FP32 staging allocations at configured capacity. +std::size_t backward_dx_staging_bytes(int max_tiles_per_reduce); +std::size_t backward_dx_configured_staging_bytes(); +std::size_t backward_dx_configured_durable_bytes(); + +void validate_backward_tp_shape( + int tokens, int hidden, int local_vocab); + +// Full resident CTA capacity used by symmetric staging on the current device. +int backward_dx_resident_cta_capacity(); + +int backward_dx_team_size(); +std::int64_t backward_dx_team_handle(); + +// The pooled dZ wave workspace for this shape. +struct BackwardScratch { + void* dz_workspace; + std::size_t dz_workspace_bytes; +}; + +BackwardScratch reserve_backward_scratch(int local_vocab); + +} // namespace fused_scaled_linear_cross_entropy +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/nvshmem.cu b/liger_cute_kernels/csrc/core/src/nvshmem.cu index eb488103..3b7c15de 100644 --- a/liger_cute_kernels/csrc/core/src/nvshmem.cu +++ b/liger_cute_kernels/csrc/core/src/nvshmem.cu @@ -24,6 +24,7 @@ #include "liger_cute/detail/comm_schedule.cuh" #include "liger_cute/detail/status.h" #include "liger_cute/detail/symmetric_memory.h" +#include "fused_scaled_linear_cross_entropy/state.h" namespace liger_cute { namespace detail { @@ -77,6 +78,7 @@ void init_comm_schedule(int N, int M) { // deterministic across PEs so the collective nvshmem_free calls stay in // lockstep. void clear_global_pools() { + liger::fused_scaled_linear_cross_entropy::reset_fslce_tp_configuration(); global_symmetric_stack().clear(); global_buffer_pool().clear(); } @@ -267,6 +269,7 @@ liger_cute_status_t liger_cute_pool_clear_all(void) { liger_cute_status_t liger_cute_pool_clear_buffers(void) { return liger_cute::detail::guarded([&]() -> liger_cute_status_t { + liger::fused_scaled_linear_cross_entropy::reset_fslce_tp_configuration(); liger_cute::detail::global_buffer_pool().clear(); return LIGER_CUTE_OK; }); diff --git a/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py b/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py index 8b75790d..535740a7 100644 --- a/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py +++ b/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py @@ -78,6 +78,28 @@ def init_pmi() -> None: _load_module().init_pmi() +def nccl_unique_id_nbytes() -> int: + out = _int64_out() + _load_module().nccl_unique_id_nbytes(out) + return int(out.item()) + + +def nccl_get_unique_id() -> torch.Tensor: + out = torch.empty(nccl_unique_id_nbytes(), dtype=torch.uint8, device="cpu") + _load_module().nccl_get_unique_id(out) + return out + + +def nccl_comm_init_rank(rank: int, nranks: int, unique_id: torch.Tensor) -> int: + out = _int64_out() + _load_module().nccl_comm_init_rank(int(rank), int(nranks), unique_id, out) + return int(out.item()) + + +def nccl_comm_destroy(comm_handle: int) -> None: + _load_module().nccl_comm_destroy(int(comm_handle)) + + def finalize() -> None: _load_module().finalize() @@ -249,3 +271,101 @@ def moe_fused_bwd_bf16( dW, ) return dX, dB, dC, dA, dW + + +def fused_linear_scaled_cross_entropy_configure_forward( + max_tokens: int, + max_local_vocab: int, +) -> None: + _load_module().fused_linear_scaled_cross_entropy_configure_forward( + int(max_tokens), + int(max_local_vocab), + ) + + +def fused_linear_scaled_cross_entropy_configure_backward( + max_tokens: int, + max_hidden: int, + max_local_vocab: int, + max_tiles_per_reduce: int, + team_handle: int, +) -> None: + _load_module().fused_linear_scaled_cross_entropy_configure_backward( + int(max_tokens), + int(max_hidden), + int(max_local_vocab), + int(max_tiles_per_reduce), + int(team_handle), + ) + + +def fused_linear_scaled_cross_entropy_forward( + x: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + vocab_start: int, + ignore_index: int, + inverse_temperature: float, + nccl_comm_handle: int, + return_entropy: bool, +): + """Run the complete tensor-parallel forward using a caller-owned NCCL communicator.""" + tokens = x.shape[0] + nll = torch.empty(tokens, dtype=torch.float32, device=x.device) + lse = torch.empty(tokens, dtype=torch.float32, device=x.device) + entropy = ( + torch.empty(tokens, dtype=torch.float32, device=x.device) + if return_entropy + else torch.zeros(tokens, dtype=torch.float32, device=x.device) + ) + _load_module().fused_linear_scaled_cross_entropy_forward( + x, + weight, + target, + int(vocab_start), + int(ignore_index), + float(inverse_temperature), + int(nccl_comm_handle), + bool(return_entropy), + nll, + lse, + entropy, + ) + return nll, lse, entropy + + +def fused_linear_scaled_cross_entropy_backward( + grad_output: torch.Tensor, + entropy_grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + lse: torch.Tensor, + entropy: torch.Tensor, + vocab_start: int, + ignore_index: int, + inverse_temperature: float, + team_handle: int, + tiles_per_reduce: int, + return_entropy: bool, +): + grad_input = torch.empty_like(x) + grad_weight = torch.empty_like(weight) + _load_module().fused_linear_scaled_cross_entropy_backward( + grad_output, + entropy_grad, + x, + weight, + target, + lse, + entropy, + int(vocab_start), + int(ignore_index), + float(inverse_temperature), + int(team_handle), + int(tiles_per_reduce), + bool(return_entropy), + grad_input, + grad_weight, + ) + return grad_input, grad_weight diff --git a/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp b/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp index 251f53c2..c3901294 100644 --- a/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp +++ b/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp @@ -2,11 +2,18 @@ #include #include +#include +#include #include #include +#include #include +#include "backward_gemm_sm90.cuh" +#include "forward_gemm_sm90.cuh" +#include "forward_reduce.cuh" +#include "workspace.cuh" #include "moe_launch.h" namespace liger { @@ -68,6 +75,13 @@ void CheckStatus(liger_cute_status_t status, const char* what) { } } +void CheckNccl(ncclResult_t status, const char* what) { + if (status != ncclSuccess) { + TVM_FFI_THROW(RuntimeError) << "liger_cute: " << what << " failed: " + << ncclGetErrorString(status); + } +} + void RequireTensor(ffi::TensorView tensor, int ndim, DLDataType dtype, const char* name) { TVM_FFI_ICHECK_EQ(tensor.ndim(), ndim) << name; TVM_FFI_ICHECK_EQ(tensor.dtype(), dtype) << name; @@ -79,6 +93,18 @@ void RequireCudaTensor(ffi::TensorView tensor, int ndim, DLDataType dtype, const TVM_FFI_ICHECK_EQ(tensor.device().device_type, kDLCUDA) << name; } +void RequireSameCudaDevice(ffi::TensorView tensor, ffi::TensorView reference, const char* name) { + TVM_FFI_ICHECK_EQ(tensor.device().device_type, reference.device().device_type) + << name << " must be on the same device as x"; + TVM_FFI_ICHECK_EQ(tensor.device().device_id, reference.device().device_id) + << name << " must be on the same device as x"; +} + +void RequirePositiveInverseTemperature(double inverse_temperature) { + TVM_FFI_ICHECK(std::isfinite(inverse_temperature) && inverse_temperature > 0.0) + << "inverse_temperature must be finite and positive"; +} + void RequireCudaMoeWeight(ffi::TensorView tensor, DLDataType dtype, const char* name) { TVM_FFI_ICHECK_EQ(tensor.ndim(), 3) << name; TVM_FFI_ICHECK_EQ(tensor.dtype(), dtype) << name; @@ -160,6 +186,52 @@ void finalize() { CheckStatus(liger_cute_nvshmem_finalize(), "finalize"); } void pool_clear_all() { CheckStatus(liger_cute_pool_clear_all(), "pool_clear_all"); } void pool_clear_buffers() { CheckStatus(liger_cute_pool_clear_buffers(), "pool_clear_buffers"); } +void nccl_unique_id_nbytes(ffi::TensorView out) { + RequireCpuInt64(out, 1); + static_cast(out.data_ptr())[0] = sizeof(ncclUniqueId); +} + +void nccl_get_unique_id(ffi::TensorView out) { + DLDataType u8{kDLUInt, 8, 1}; + RequireTensor(out, 1, u8, "out"); + TVM_FFI_ICHECK_EQ(out.device().device_type, kDLCPU); + TVM_FFI_ICHECK_EQ(out.size(0), static_cast(sizeof(ncclUniqueId))); + ncclUniqueId unique_id; + CheckNccl(ncclGetUniqueId(&unique_id), "ncclGetUniqueId"); + std::memcpy(out.data_ptr(), &unique_id, sizeof(unique_id)); +} + +void nccl_comm_init_rank( + int64_t rank, int64_t nranks, ffi::TensorView unique_id, ffi::TensorView out) { + DLDataType u8{kDLUInt, 8, 1}; + RequireTensor(unique_id, 1, u8, "unique_id"); + TVM_FFI_ICHECK_EQ(unique_id.device().device_type, kDLCPU); + TVM_FFI_ICHECK_EQ( + unique_id.size(0), static_cast(sizeof(ncclUniqueId))); + RequireCpuInt64(out, 1); + TVM_FFI_ICHECK_GT(nranks, 0); + TVM_FFI_ICHECK_GE(rank, 0); + TVM_FFI_ICHECK_LT(rank, nranks); + ncclUniqueId id; + std::memcpy(&id, unique_id.data_ptr(), sizeof(id)); + ncclComm_t comm = nullptr; + CheckNccl( + ncclCommInitRank( + &comm, static_cast(nranks), id, static_cast(rank)), + "ncclCommInitRank"); + static_cast(out.data_ptr())[0] = + reinterpret_cast(comm); +} + +void nccl_comm_destroy(int64_t comm_handle) { + if (comm_handle == 0) { + return; + } + CheckNccl( + ncclCommDestroy(reinterpret_cast(comm_handle)), + "ncclCommDestroy"); +} + void my_pe(ffi::TensorView out) { RequireCpuInt32(out, 1); CheckStatus(liger_cute_nvshmem_my_pe(static_cast(out.data_ptr())), "my_pe"); @@ -384,6 +456,181 @@ void moe_fused_bwd_bf16( } } +void fused_linear_scaled_cross_entropy_configure_forward( + int64_t max_tokens, int64_t max_local_vocab) { + TVM_FFI_ICHECK_GT(max_tokens, 0); + TVM_FFI_ICHECK_GT(max_local_vocab, 0); + try { + liger::fused_scaled_linear_cross_entropy::configure_forward_tp_workspace( + static_cast(max_tokens), static_cast(max_local_vocab)); + } catch (const std::exception& e) { + ThrowCoreError("fused_linear_scaled_cross_entropy_configure_forward", e); + } +} + +void fused_linear_scaled_cross_entropy_configure_backward( + int64_t max_tokens, int64_t max_hidden, int64_t max_local_vocab, + int64_t max_tiles_per_reduce, int64_t team_handle) { + TVM_FFI_ICHECK_GT(max_tokens, 0); + TVM_FFI_ICHECK_GT(max_hidden, 0); + TVM_FFI_ICHECK_GT(max_local_vocab, 0); + TVM_FFI_ICHECK( + max_tiles_per_reduce == 1 || max_tiles_per_reduce == 2 || + max_tiles_per_reduce == 4); + try { + liger::fused_scaled_linear_cross_entropy::configure_backward_tp_symmetric( + static_cast(max_tokens), static_cast(max_hidden), + static_cast(max_local_vocab), static_cast(max_tiles_per_reduce), 1, team_handle); + } catch (const std::exception& e) { + ThrowCoreError("fused_linear_scaled_cross_entropy_configure_backward", e); + } +} + +void fused_linear_scaled_cross_entropy_forward( + ffi::TensorView x, ffi::TensorView weight, ffi::TensorView target, + int64_t vocab_start, int64_t ignore_index, double inverse_temperature, + int64_t nccl_comm_handle, bool return_entropy, ffi::TensorView nll, + ffi::TensorView lse, ffi::TensorView entropy) { + DLDataType bf16{kDLBfloat, 16, 1}; + DLDataType f32{kDLFloat, 32, 1}; + DLDataType i64{kDLInt, 64, 1}; + RequireCudaTensor(x, 2, bf16, "x"); + RequireCudaTensor(weight, 2, bf16, "weight"); + RequireCudaTensor(target, 1, i64, "target"); + RequireCudaTensor(nll, 1, f32, "nll"); + RequireCudaTensor(lse, 1, f32, "lse"); + RequireCudaTensor(entropy, 1, f32, "entropy"); + RequireSameCudaDevice(weight, x, "weight"); + RequireSameCudaDevice(target, x, "target"); + RequireSameCudaDevice(nll, x, "nll"); + RequireSameCudaDevice(lse, x, "lse"); + RequireSameCudaDevice(entropy, x, "entropy"); + + int64_t tokens = x.size(0); + int64_t hidden = x.size(1); + int64_t local_vocab = weight.size(0); + TVM_FFI_ICHECK_EQ(weight.size(1), hidden); + TVM_FFI_ICHECK_EQ(target.size(0), tokens); + TVM_FFI_ICHECK_EQ(nll.size(0), tokens); + TVM_FFI_ICHECK_EQ(lse.size(0), tokens); + TVM_FFI_ICHECK_EQ(entropy.size(0), tokens); + TVM_FFI_ICHECK_GE(vocab_start, 0); + TVM_FFI_ICHECK_NE(nccl_comm_handle, 0); + RequirePositiveInverseTemperature(inverse_temperature); + + using namespace liger::fused_scaled_linear_cross_entropy; + ForwardTpParamsSm90<90> params; + params.gemm.x = x.data_ptr(); + params.gemm.weight = weight.data_ptr(); + params.gemm.target = static_cast(target.data_ptr()); + params.gemm.tokens = static_cast(tokens); + params.gemm.hidden = static_cast(hidden); + params.gemm.local_vocab = static_cast(local_vocab); + params.gemm.vocab_start = vocab_start; + params.gemm.ignore_index = ignore_index; + params.gemm.inverse_temperature = static_cast(inverse_temperature); + params.nll = static_cast(nll.data_ptr()); + params.lse = static_cast(lse.data_ptr()); + params.entropy = static_cast(entropy.data_ptr()); + params.nccl_comm_handle = nccl_comm_handle; + cudaStream_t stream = reinterpret_cast( + TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); + try { + if (return_entropy) { + liger::fused_scaled_linear_cross_entropy:: + fused_linear_scaled_cross_entropy_forward(params, stream); + } else { + liger::fused_scaled_linear_cross_entropy:: + fused_linear_scaled_cross_entropy_forward(params, stream); + } + } catch (const std::exception& e) { + ThrowCoreError("fused_linear_scaled_cross_entropy_forward", e); + } +} + +void fused_linear_scaled_cross_entropy_backward( + ffi::TensorView grad_output, ffi::TensorView entropy_grad, ffi::TensorView x, + ffi::TensorView weight, ffi::TensorView target, ffi::TensorView lse, + ffi::TensorView entropy, int64_t vocab_start, int64_t ignore_index, + double inverse_temperature, int64_t team_handle, int64_t tiles_per_reduce, + bool return_entropy, ffi::TensorView grad_input, ffi::TensorView grad_weight) { + DLDataType bf16{kDLBfloat, 16, 1}; + DLDataType f32{kDLFloat, 32, 1}; + DLDataType i64{kDLInt, 64, 1}; + RequireCudaTensor(grad_output, 1, f32, "grad_output"); + RequireCudaTensor(entropy_grad, 1, f32, "entropy_grad"); + RequireCudaTensor(x, 2, bf16, "x"); + RequireCudaTensor(weight, 2, bf16, "weight"); + RequireCudaTensor(target, 1, i64, "target"); + RequireCudaTensor(lse, 1, f32, "lse"); + RequireCudaTensor(entropy, 1, f32, "entropy"); + RequireCudaTensor(grad_input, 2, bf16, "grad_input"); + RequireCudaTensor(grad_weight, 2, bf16, "grad_weight"); + RequireSameCudaDevice(entropy_grad, x, "entropy_grad"); + RequireSameCudaDevice(weight, x, "weight"); + RequireSameCudaDevice(target, x, "target"); + RequireSameCudaDevice(grad_output, x, "grad_output"); + RequireSameCudaDevice(lse, x, "lse"); + RequireSameCudaDevice(entropy, x, "entropy"); + RequireSameCudaDevice(grad_input, x, "grad_input"); + RequireSameCudaDevice(grad_weight, x, "grad_weight"); + + int64_t tokens = x.size(0); + int64_t hidden = x.size(1); + int64_t local_vocab = weight.size(0); + TVM_FFI_ICHECK_EQ(weight.size(1), hidden); + TVM_FFI_ICHECK_EQ(target.size(0), tokens); + TVM_FFI_ICHECK_EQ(grad_output.size(0), tokens); + TVM_FFI_ICHECK_EQ(entropy_grad.size(0), tokens); + TVM_FFI_ICHECK_EQ(lse.size(0), tokens); + TVM_FFI_ICHECK_EQ(entropy.size(0), tokens); + TVM_FFI_ICHECK_EQ(grad_input.size(0), tokens); + TVM_FFI_ICHECK_EQ(grad_input.size(1), hidden); + TVM_FFI_ICHECK_EQ(grad_weight.size(0), local_vocab); + TVM_FFI_ICHECK_EQ(grad_weight.size(1), hidden); + TVM_FFI_ICHECK( + tiles_per_reduce == 1 || tiles_per_reduce == 2 || tiles_per_reduce == 4); + RequirePositiveInverseTemperature(inverse_temperature); + + using namespace liger::fused_scaled_linear_cross_entropy; + cudaStream_t stream = reinterpret_cast( + TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); + try { + TVM_FFI_ICHECK_EQ(team_handle, backward_dx_team_handle()) + << "team_handle must match the configured tensor-parallel team"; + BackwardScratch scratch = reserve_backward_scratch(static_cast(local_vocab)); + BackwardTpParamsSm90<90> params; + params.gemm.x = x.data_ptr(); + params.gemm.weight = weight.data_ptr(); + params.gemm.target = static_cast(target.data_ptr()); + params.gemm.grad_output = static_cast(grad_output.data_ptr()); + params.gemm.lse = static_cast(lse.data_ptr()); + params.gemm.entropy = static_cast(entropy.data_ptr()); + params.gemm.entropy_grad = static_cast(entropy_grad.data_ptr()); + params.gemm.grad_input = grad_input.data_ptr(); + params.gemm.grad_weight = grad_weight.data_ptr(); + params.gemm.dz_workspace = scratch.dz_workspace; + params.gemm.dz_workspace_bytes = scratch.dz_workspace_bytes; + params.gemm.tokens = static_cast(tokens); + params.gemm.hidden = static_cast(hidden); + params.gemm.local_vocab = static_cast(local_vocab); + params.gemm.vocab_start = vocab_start; + params.gemm.ignore_index = ignore_index; + params.gemm.inverse_temperature = static_cast(inverse_temperature); + params.tiles_per_reduce = static_cast(tiles_per_reduce); + params.num_comm_channels = 1; + if (return_entropy) { + liger::fused_scaled_linear_cross_entropy:: + fused_linear_scaled_cross_entropy_backward(params, stream); + } else { + liger::fused_scaled_linear_cross_entropy:: + fused_linear_scaled_cross_entropy_backward(params, stream); + } + } catch (const std::exception& e) { + ThrowCoreError("fused_linear_scaled_cross_entropy_backward", e); + } +} + } // namespace TVM_FFI_DLL_EXPORT_TYPED_FUNC(uniqueid_nbytes, uniqueid_nbytes); @@ -391,6 +638,10 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_uniqueid, get_uniqueid); TVM_FFI_DLL_EXPORT_TYPED_FUNC(init_with_uniqueid, init_with_uniqueid); TVM_FFI_DLL_EXPORT_TYPED_FUNC(init_pmi, init_pmi); TVM_FFI_DLL_EXPORT_TYPED_FUNC(finalize, finalize); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(nccl_unique_id_nbytes, nccl_unique_id_nbytes); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(nccl_get_unique_id, nccl_get_unique_id); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(nccl_comm_init_rank, nccl_comm_init_rank); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(nccl_comm_destroy, nccl_comm_destroy); TVM_FFI_DLL_EXPORT_TYPED_FUNC(my_pe, my_pe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(n_pes, n_pes); TVM_FFI_DLL_EXPORT_TYPED_FUNC(team_world, team_world); @@ -406,3 +657,15 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_configure_symmetric, moe_configure_symmetric); TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_pop_fwd, moe_pop_fwd); TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_fused_fwd_bf16, moe_fused_fwd_bf16); TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_fused_bwd_bf16, moe_fused_bwd_bf16); +TVM_FFI_DLL_EXPORT_TYPED_FUNC( + fused_linear_scaled_cross_entropy_configure_forward, + fused_linear_scaled_cross_entropy_configure_forward); +TVM_FFI_DLL_EXPORT_TYPED_FUNC( + fused_linear_scaled_cross_entropy_configure_backward, + fused_linear_scaled_cross_entropy_configure_backward); +TVM_FFI_DLL_EXPORT_TYPED_FUNC( + fused_linear_scaled_cross_entropy_forward, + fused_linear_scaled_cross_entropy_forward); +TVM_FFI_DLL_EXPORT_TYPED_FUNC( + fused_linear_scaled_cross_entropy_backward, + fused_linear_scaled_cross_entropy_backward); diff --git a/liger_cute_kernels/test/test_fslce_teams.py b/liger_cute_kernels/test/test_fslce_teams.py new file mode 100644 index 00000000..aa259e2a --- /dev/null +++ b/liger_cute_kernels/test/test_fslce_teams.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +import os +import shutil +import sys +import tempfile + +from datetime import timedelta + +import pytest + +try: + import torch + import torch.distributed as dist + import torch.multiprocessing as mp + + from liger_cute_kernels import nvshmem + from liger_cute_kernels import tvm_ffi +except ImportError: + torch = None + dist = None + mp = None + nvshmem = None + tvm_ffi = None + +pytestmark = pytest.mark.skipif( + nvshmem is None, + reason="liger_cute_kernels dependencies are not installed", +) + +_NATIVE_AVAILABLE = tvm_ffi is not None and tvm_ffi.is_available() +_NDEV = torch.cuda.device_count() if _NATIVE_AVAILABLE and torch is not None and torch.cuda.is_available() else 0 + + +def test_fslce_forward_binding_is_exposed(): + if not _NATIVE_AVAILABLE: + pytest.skip("requires the native core") + module = tvm_ffi._load_module() + for name in ( + "nccl_get_unique_id", + "nccl_comm_init_rank", + "fused_linear_scaled_cross_entropy_configure_forward", + "fused_linear_scaled_cross_entropy_forward", + ): + assert hasattr(module, name), f"missing native TVM FFI export: {name}" + + +def test_forward_returns_zero_entropy_when_disabled(monkeypatch): + class FakeModule: + @staticmethod + def fused_linear_scaled_cross_entropy_forward(*args): + del args + + monkeypatch.setattr(tvm_ffi, "_load_module", lambda: FakeModule()) + x = torch.empty((2, 8), dtype=torch.bfloat16) + weight = torch.empty((4, 8), dtype=torch.bfloat16) + target = torch.zeros(2, dtype=torch.int64) + _, _, entropy = tvm_ffi.fused_linear_scaled_cross_entropy_forward( + x, + weight, + target, + vocab_start=0, + ignore_index=-100, + inverse_temperature=1.0, + nccl_comm_handle=1, + return_entropy=False, + ) + torch.testing.assert_close(entropy, torch.zeros_like(entropy)) + + +def _reference(x, weights, target, grad_output, temperature): + logits = (x.float() @ weights.float().T) / temperature + lse = torch.logsumexp(logits, dim=-1) + log_probs = logits - lse[:, None] + probabilities = torch.softmax(logits, dim=-1) + rows = torch.arange(x.shape[0], device=x.device) + nll = -log_probs[rows, target] + entropy = -(probabilities * log_probs).sum(dim=-1) + dz = probabilities + dz[torch.arange(x.shape[0], device=x.device), target] -= 1.0 + dz *= grad_output[:, None] / temperature + dz = dz.to(torch.bfloat16).float() + return nll, lse, entropy, dz @ weights.float(), dz + + +def _worker(rank: int, world_size: int, init_file: str, subgroup: bool): + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=180), + ) + + team = None + nccl_comm = 0 + try: + nvshmem.init_from_pg() + pg = dist.group.WORLD + group_ranks = list(range(world_size)) + if subgroup: + groups = ([0, 2], [1, 3]) + for ranks in groups: + candidate = dist.new_group(ranks=ranks) + if rank in ranks: + pg = candidate + group_ranks = list(ranks) + team = nvshmem.resolve_team(pg) + else: + team = nvshmem.resolve_team(None) + + local_rank = group_ranks.index(rank) + team_size = len(group_ranks) + unique_id_payload = [tvm_ffi.nccl_get_unique_id().tolist() if local_rank == 0 else None] + dist.broadcast_object_list( + unique_id_payload, + src=group_ranks[0], + group=pg, + device=torch.device("cuda", rank), + ) + unique_id = torch.tensor(unique_id_payload[0], dtype=torch.uint8) + nccl_comm = tvm_ffi.nccl_comm_init_rank(local_rank, team_size, unique_id) + + tokens = 128 + hidden = 2048 + local_vocab = 320 + temperature = 0.8 + global_vocab = team_size * local_vocab + + generator = torch.Generator(device="cpu") + generator.manual_seed(2027 + (rank & 1 if subgroup else 0)) + x = torch.randn(tokens, hidden, generator=generator).mul_(0.05).to(device="cuda", dtype=torch.bfloat16) + weight_generator = torch.Generator(device="cpu") + weight_generator.manual_seed(3100 + local_rank) + weight = ( + torch.randn(local_vocab, hidden, generator=weight_generator) + .mul_(0.05) + .to(device="cuda", dtype=torch.bfloat16) + ) + gathered = [torch.empty_like(weight) for _ in range(team_size)] + dist.all_gather(gathered, weight, group=pg) + global_weight = torch.cat(gathered, dim=0) + + target = torch.arange(tokens, device="cuda", dtype=torch.int64) % global_vocab + grad_output = torch.linspace(0.5, 1.0, tokens, device="cuda", dtype=torch.float32) + entropy = torch.zeros(tokens, device="cuda", dtype=torch.float32) + entropy_grad = torch.zeros_like(entropy) + expected_nll, expected_lse, expected_entropy, expected_dx, dz = _reference( + x, global_weight, target, grad_output, temperature + ) + expected_dw = dz[:, local_rank * local_vocab : (local_rank + 1) * local_vocab].T @ x.float() + + tvm_ffi.fused_linear_scaled_cross_entropy_configure_forward(tokens, local_vocab) + tvm_ffi.fused_linear_scaled_cross_entropy_configure_backward(tokens, hidden, local_vocab, 1, team) + actual_nll, lse, entropy = tvm_ffi.fused_linear_scaled_cross_entropy_forward( + x, + weight, + target, + local_rank * local_vocab, + -100, + 1.0 / temperature, + nccl_comm, + True, + ) + actual_dx, actual_dw = tvm_ffi.fused_linear_scaled_cross_entropy_backward( + grad_output, + entropy_grad, + x, + weight, + target, + lse, + entropy, + local_rank * local_vocab, + -100, + 1.0 / temperature, + team, + 1, + False, + ) + torch.cuda.synchronize() + torch.testing.assert_close(actual_nll, expected_nll, atol=2e-4, rtol=2e-4) + torch.testing.assert_close(lse, expected_lse, atol=2e-4, rtol=2e-4) + torch.testing.assert_close(entropy, expected_entropy, atol=2e-4, rtol=2e-4) + torch.testing.assert_close(actual_dx.float(), expected_dx, atol=8e-3, rtol=4e-2) + torch.testing.assert_close(actual_dw.float(), expected_dw, atol=8e-3, rtol=4e-2) + + nvshmem.pool_clear_all() + tvm_ffi.fused_linear_scaled_cross_entropy_configure_backward(tokens, hidden, local_vocab, 1, team) + replay_dx = torch.empty_like(x) + replay_dw = torch.empty_like(weight) + module = tvm_ffi._load_module() + graph = torch.cuda.CUDAGraph() + dist.barrier() + with torch.cuda.graph(graph): + module.fused_linear_scaled_cross_entropy_backward( + grad_output, + entropy_grad, + x, + weight, + target, + lse, + entropy, + local_rank * local_vocab, + -100, + 1.0 / temperature, + team, + 1, + False, + replay_dx, + replay_dw, + ) + dist.barrier() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(replay_dx.float(), expected_dx, atol=8e-3, rtol=4e-2) + torch.testing.assert_close(replay_dw.float(), expected_dw, atol=8e-3, rtol=4e-2) + replay_dx.fill_(float("nan")) + replay_dw.fill_(float("nan")) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(replay_dx.float(), expected_dx, atol=8e-3, rtol=4e-2) + torch.testing.assert_close(replay_dw.float(), expected_dw, atol=8e-3, rtol=4e-2) + + tvm_ffi.nccl_comm_destroy(nccl_comm) + nccl_comm = 0 + nvshmem.pool_clear_all() + if subgroup and team != nvshmem.team_world(): + nvshmem.team_destroy(team) + team = None + dist.barrier() + nvshmem.finalize() + dist.destroy_process_group() + except BaseException: + if nccl_comm: + try: + tvm_ffi.nccl_comm_destroy(nccl_comm) + except Exception as cleanup_error: + print(f"NCCL communicator cleanup failed: {cleanup_error}", file=sys.stderr) + try: + dist.destroy_process_group() + except Exception: + pass + raise + + +def _run(world_size: int, subgroup: bool): + rendezvous = tempfile.mkdtemp(prefix="fslce_team_") + init_file = os.path.join(rendezvous, "store") + old_heap = os.environ.get("NVSHMEM_SYMMETRIC_SIZE") + os.environ["NVSHMEM_SYMMETRIC_SIZE"] = "3G" + try: + mp.spawn( + _worker, + args=(world_size, init_file, subgroup), + nprocs=world_size, + join=True, + ) + finally: + if old_heap is None: + os.environ.pop("NVSHMEM_SYMMETRIC_SIZE", None) + else: + os.environ["NVSHMEM_SYMMETRIC_SIZE"] = old_heap + shutil.rmtree(rendezvous, ignore_errors=True) + + +@pytest.mark.skipif( + not _NATIVE_AVAILABLE or _NDEV < 2, + reason="requires the native core and at least two GPUs", +) +def test_fslce_forward_backward_world_team(): + _run(2, subgroup=False) + + +@pytest.mark.skipif( + not _NATIVE_AVAILABLE or _NDEV < 4, + reason="requires the native core and at least four GPUs", +) +def test_fslce_forward_backward_process_group_subteams(): + _run(4, subgroup=True) diff --git a/liger_cute_kernels/tests/cpp/CMakeLists.txt b/liger_cute_kernels/tests/cpp/CMakeLists.txt index fd71ad40..8a993539 100644 --- a/liger_cute_kernels/tests/cpp/CMakeLists.txt +++ b/liger_cute_kernels/tests/cpp/CMakeLists.txt @@ -27,6 +27,20 @@ target_include_directories(test_mlp1_fused PRIVATE "${CMAKE_SOURCE_DIR}/csrc/core/src/moe") gtest_discover_tests(test_mlp1_fused DISCOVERY_MODE PRE_TEST) +add_executable( + test_fused_scaled_linear_cross_entropy + test_fused_scaled_linear_cross_entropy.cu) +target_link_libraries(test_fused_scaled_linear_cross_entropy PRIVATE + GTest::gtest + CUTLASS::CUTLASS + CUDA::cudart) +target_include_directories(test_fused_scaled_linear_cross_entropy PRIVATE + "${CMAKE_SOURCE_DIR}/csrc/core/include" + "${CMAKE_SOURCE_DIR}/csrc/core/src/fused_scaled_linear_cross_entropy") +gtest_discover_tests( + test_fused_scaled_linear_cross_entropy + DISCOVERY_MODE PRE_TEST) + # ── Blackwell (SM100a) port targets — pre-staged by the orchestrator so the # three parallel port subagents each build ONLY their own already-registered # target and never edit this shared file. Each test .cu starts as a minimal stub diff --git a/liger_cute_kernels/tests/cpp/test_fused_scaled_linear_cross_entropy.cu b/liger_cute_kernels/tests/cpp/test_fused_scaled_linear_cross_entropy.cu new file mode 100644 index 00000000..15b39762 --- /dev/null +++ b/liger_cute_kernels/tests/cpp/test_fused_scaled_linear_cross_entropy.cu @@ -0,0 +1,195 @@ +#include + +#include + +#include +#include + +#include "backward_gemm_mainloop_sm90.cuh" +#include "fused_scaled_linear_cross_entropy.cuh" + +namespace fslce = liger::fused_scaled_linear_cross_entropy; + +using ForwardConfig = fslce::ForwardGemmConfigSm90<90>; +using BackwardConfig = fslce::BackwardGemmConfigSm90<90>; +using CommConfig = fslce::DxCommConfig; + +static_assert(ForwardConfig::kCompute == 90); +static_assert(BackwardConfig::kCompute == 90); +static_assert(fslce::BackwardDzHandoffTraitsSm90<90>::kStages == 3); +static_assert(fslce::BackwardDxCluster2TraitsSm90<90>::kStages == 4); +static_assert(fslce::BackwardDxCluster2TraitsSm90<90>::kStoreTileN == 32); +static_assert(sizeof(fslce::BackwardDxCluster2SmemSm90<90>) == 230400); +static_assert(sizeof(fslce::DxReduceWorkspace) == 4 * sizeof(void*)); +static_assert(sizeof(liger_cute::detail::NvlsReduceView) <= 5 * sizeof(void*)); +static_assert( + sizeof(liger_cute::detail::DirectPeerReduceView) <= 4 * sizeof(void*)); +static_assert(fslce::backward_dx_split_k(2048, 256) == 2); +static_assert(fslce::backward_dx_split_k(2048, 320) == 1); +static_assert(fslce::backward_dx_split_k(4096, 256) == 1); + +TEST(TensorParallelFusedScaledLinearCrossEntropyConfig, PreservesHopperWarpRoles) { + EXPECT_EQ( + fslce::backward_warp_role(0), + fslce::BackwardWarpRole::kProducer); + EXPECT_EQ( + fslce::backward_warp_role(1), + fslce::BackwardWarpRole::kDxCommunication); + EXPECT_EQ( + fslce::backward_warp_role(2), + fslce::BackwardWarpRole::kDxCommunication); + EXPECT_EQ( + fslce::backward_warp_role(3), + fslce::BackwardWarpRole::kReserved); + for (int warp = 4; warp < fslce::kNumWarps; ++warp) { + EXPECT_EQ( + fslce::backward_warp_role(warp), + fslce::BackwardWarpRole::kConsumer); + } + EXPECT_EQ( + fslce::backward_warp_role(fslce::kNumWarps), + fslce::BackwardWarpRole::kInactive); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyConfig, CoalescesTwoDxTilesByDefault) { + EXPECT_EQ(fslce::dx_groups_per_m_tile(5), 3); + + fslce::DxTileGroup first = + fslce::make_dx_tile_group(3, 0, 5, 7); + EXPECT_EQ(first.m_tile, 3); + EXPECT_EQ(first.first_n_tile, 0); + EXPECT_EQ(first.num_tiles, 2); + EXPECT_EQ(first.sequence, 7); + + fslce::DxTileGroup tail = + fslce::make_dx_tile_group(3, 2, 5, 9); + EXPECT_EQ(tail.first_n_tile, 4); + EXPECT_EQ(tail.num_tiles, 1); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyOnlineSoftmax, MatchesGlobalReference) { + fslce::OnlineSoftmaxState rank0; + rank0.fold(1.0f, true, false); + rank0.fold(3.0f, true, false); + + fslce::OnlineSoftmaxState rank1; + rank1.fold(2.0f, true, true); + rank1.fold(4.0f, true, false); + + float global_max = fmaxf(rank0.max_value, rank1.max_value); + fslce::CorrectedSoftmaxStats corrected0 = + fslce::correct_for_global_max(rank0, global_max); + fslce::CorrectedSoftmaxStats corrected1 = + fslce::correct_for_global_max(rank1, global_max); + + float global_sum = corrected0.exp_sum + corrected1.exp_sum; + float global_weighted = + corrected0.exp_weighted_sum + corrected1.exp_weighted_sum; + float global_target = + corrected0.target_logit + corrected1.target_logit; + fslce::FinalizedSoftmax result = fslce::finalize_softmax( + global_max, + global_sum, + global_target, + global_weighted, + false); + + float direct_sum = + expf(1.0f - global_max) + + expf(2.0f - global_max) + + expf(3.0f - global_max) + + expf(4.0f - global_max); + float expected_lse = global_max + logf(direct_sum); + float direct_weighted = + expf(1.0f - global_max) * 1.0f + + expf(2.0f - global_max) * 2.0f + + expf(3.0f - global_max) * 3.0f + + expf(4.0f - global_max) * 4.0f; + + EXPECT_NEAR(result.lse, expected_lse, 1.0e-6f); + EXPECT_NEAR(result.nll, expected_lse - 2.0f, 1.0e-6f); + EXPECT_NEAR( + result.entropy, + expected_lse - direct_weighted / direct_sum, + 1.0e-6f); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyOnlineSoftmax, ChunkMergeMatchesSingleFold) { + fslce::OnlineSoftmaxState lhs; + lhs.fold(-2.0f, true, false); + lhs.fold(0.5f, true, false); + + fslce::OnlineSoftmaxState rhs; + rhs.fold(3.0f, true, true); + rhs.fold(1.0f, true, false); + + fslce::OnlineSoftmaxState merged = + fslce::merge_online_softmax(lhs, rhs); + fslce::OnlineSoftmaxState direct; + direct.fold(-2.0f, true, false); + direct.fold(0.5f, true, false); + direct.fold(3.0f, true, true); + direct.fold(1.0f, true, false); + + EXPECT_NEAR(merged.max_value, direct.max_value, 1.0e-7f); + EXPECT_NEAR(merged.exp_sum, direct.exp_sum, 1.0e-6f); + EXPECT_NEAR( + merged.exp_weighted_sum, + direct.exp_weighted_sum, + 1.0e-6f); + EXPECT_EQ(merged.has_target, 1); + EXPECT_FLOAT_EQ(merged.target_logit, 3.0f); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyOnlineSoftmax, CompilesOutEntropyMoment) { + fslce::OnlineSoftmaxState state; + state.fold(2.0f, true, false); + state.fold(4.0f, true, false); + EXPECT_FLOAT_EQ(state.exp_weighted_sum, 0.0f); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyOnlineSoftmax, ZerosIgnoredEntropy) { + fslce::FinalizedSoftmax result = fslce::finalize_softmax( + 4.0f, + 2.0f, + 0.0f, + 6.0f, + true); + EXPECT_FLOAT_EQ(result.nll, 0.0f); + EXPECT_FLOAT_EQ(result.entropy, 0.0f); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyBackward, AppliesTemperatureChainRule) { + float temperature = 2.0f; + float scaled_logit = 0.25f; + float lse = 1.25f; + float scale = 0.75f; + float probability = expf(scaled_logit - lse); + + float gradient = fslce::scaled_cross_entropy_dlogit( + scaled_logit, + lse, + true, + scale, + temperature); + float expected = scale * (probability - 1.0f) / temperature; + EXPECT_NEAR(gradient, expected, 1.0e-7f); +} + +TEST(TensorParallelFusedScaledLinearCrossEntropyBackward, ZerosIgnoredRows) { + float gradient = fslce::scaled_cross_entropy_dlogit( + 0.25f, + 1.25f, + true, + 1.0f, + 2.0f, + 0.5f, + 1.0f, + true); + EXPECT_FLOAT_EQ(gradient, 0.0f); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}