linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CUDA] Add Hopper tensor-parallel fused linear cross entropy (#1422)

Browse Source 
## 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 <[email protected]>
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Jari Kolehmainen
2026-08-31 20:15:17 +00:00
committed by GitHub
co-authored by jkolehma Copilot
parent f65ae6576f
commit 99747c53ba
38 changed files with 10731 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
liger_cute_kernels/CMakeLists.txt
+18 -4
View File
@@ -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=<prefix>)")
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 ──────────────────────────────────────────────────────────────────
liger_cute_kernels/README.md
+35
View File
@@ -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
liger_cute_kernels/csrc/core/CMakeLists.txt
+106 -1
View File
@@ -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
$<$<COMPILE_LANGUAGE:CUDA>:-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}
$<TARGET_OBJECTS:liger_cute_fslce_forward_sm90>
$<TARGET_OBJECTS:liger_cute_fslce_forward_reduce>
$<TARGET_OBJECTS:liger_cute_fslce_backward_sm90>
$<TARGET_OBJECTS:liger_cute_fslce_backward_reduce_sm90>
"${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
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
$<INSTALL_INTERFACE:include>
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
)
liger_cute_kernels/csrc/core/include/liger_cute/detail/local_reduce.cuh
+154
View File
@@ -0,0 +1,154 @@
#pragma once
// Device-side local reduction algorithms for TpReducePlan backends.
#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include "liger_cute/detail/nvls.cuh"
#include "liger_cute/detail/tp_reduce.cuh"
namespace liger_cute {
namespace detail {
template <LocalReduceBackend Backend, typename Element>
struct LocalReduceContext;
template <typename Element>
struct LocalReduceContext<LocalReduceBackend::kNvls, Element> {
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 <typename Element>
struct LocalReduceContext<LocalReduceBackend::kDirectPeer, Element> {
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<int>(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<int>(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 <LocalReduceBackend Backend, typename Element>
__device__ __forceinline__ void local_all_reduce(
const LocalReduceContext<Backend, Element>& 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<Element, float>,
"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<Element, float>,
"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<const float4*>(
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<float4*>(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
liger_cute_kernels/csrc/core/include/liger_cute/detail/nvls.cuh
+304
View File
@@ -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 <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
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<int>(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<std::size_t>(team_size);
const std::size_t rank_offset =
elements_per_pe * static_cast<std::size_t>(team_rank);
const std::size_t local_count =
elements_per_pe +
(team_rank == team_size - 1
? count % static_cast<std::size_t>(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<std::uintptr_t>(source) & 0xF) == 0 &&
(reinterpret_cast<std::uintptr_t>(dest) & 0xF) == 0) {
const std::size_t vectors = remaining / 4;
const auto* source_vectors = reinterpret_cast<const int4*>(source);
auto* dest_vectors = reinterpret_cast<int4*>(dest);
for (std::size_t index = static_cast<std::size_t>(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<std::uintptr_t>(source) & 0x7) == 0 &&
(reinterpret_cast<std::uintptr_t>(dest) & 0x7) == 0) {
const std::size_t vectors = remaining / 2;
const auto* source_vectors =
reinterpret_cast<const std::uint64_t*>(source);
auto* dest_vectors = reinterpret_cast<std::uint64_t*>(dest);
for (std::size_t index = static_cast<std::size_t>(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<std::size_t>(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<std::size_t>(size);
return per_rank +
(rank == size - 1 ? count % static_cast<std::size_t>(size) : 0);
}
__device__ __forceinline__ std::size_t nvls_rank_offset(
std::size_t count, int rank, int size) {
return count / static_cast<std::size_t>(size) *
static_cast<std::size_t>(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<const int4*>(source) + index)
: "memory");
asm volatile(
"st.global.v4.b32 [%0], {%1, %2, %3, %4};"
:
: "l"(reinterpret_cast<int4*>(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<const int4*>(local_source)[index];
asm volatile(
"multimem.st.relaxed.sys.global.v4.f32 "
"[%0], {%1, %2, %3, %4};"
:
: "l"(reinterpret_cast<int4*>(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<const std::uint32_t*>(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
liger_cute_kernels/csrc/core/include/liger_cute/detail/tp_reduce.cuh
+111
View File
@@ -0,0 +1,111 @@
#pragma once
#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
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
liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cu
+84
View File
@@ -0,0 +1,84 @@
#include "remote_all_reduce.cuh"
#include <nvshmem.h>
#include <nvshmemx.h>
#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<int>(
(count + 255) / 256 < 1024 ? (count + 255) / 256 : 1024);
remote_sum<<<blocks, 256, 0, stream>>>(
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
liger_cute_kernels/csrc/core/src/detail/remote_all_reduce.cuh
+26
View File
@@ -0,0 +1,26 @@
#pragma once
// Internal implementation detail behind launch_remote_reduce().
#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
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
liger_cute_kernels/csrc/core/src/detail/tp_reduce.cu
+357
View File
@@ -0,0 +1,357 @@
#include "liger_cute/detail/tp_reduce.cuh"
#include <cuda_runtime.h>
#include <nvshmem.h>
#include <nvshmemx.h>
#include <cstddef>
#include <cstdint>
#include <vector>
#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<float*>(nvshmemx_mc_ptr(team, partial));
output->multicast_reduced =
static_cast<float*>(nvshmemx_mc_ptr(team, reduced));
output->multicast_sync =
static_cast<std::uint64_t*>(nvshmemx_mc_ptr(team, sync));
output->node_multicast_partial =
static_cast<float*>(
nvshmemx_mc_ptr(NVSHMEMX_TEAM_NODE, partial));
output->node_multicast_reduced =
static_cast<float*>(
nvshmemx_mc_ptr(NVSHMEMX_TEAM_NODE, reduced));
output->node_multicast_sync =
static_cast<std::uint64_t*>(
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<nvshmem_team_t>(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<nvshmem_team_t>(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<float*> peer_partial(team_size);
std::vector<std::uint64_t*> peer_sync(team_size);
std::vector<unsigned char> world_members(
static_cast<std::size_t>(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<int>(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<float*>(
nvshmem_ptr(buffers.partial, world_pe));
peer_sync[rank] = world_pe == my_world_pe
? buffers.sync
: static_cast<std::uint64_t*>(
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<int>(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<std::uint64_t*>(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<nvshmem_team_t>(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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward.cuh
+47
View File
@@ -0,0 +1,47 @@
#pragma once
#include <cmath>
#include <cuda_runtime.h>
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 <typename GemmConfig>
__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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm90.cuh
+1907
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_sm90.cuh
+381
View File
@@ -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 <cuda_runtime.h>
#include <cmath>
#include <cstddef>
#include <cstdint>
#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 <int Compute = 90>
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 <int Compute = 90>
struct DzGemmContractSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct DxGemmContractSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
static constexpr int kCompute = Compute;
static constexpr int kTileM = Config::kTileM;
static constexpr int kTileN = Config::kDxTileN;
static constexpr int kTileK = Config::kTileK;
};
template <int Compute = 90>
struct DwGemmContractSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct BackwardWarpPlanSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct BackwardGemmParamsSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct BackwardGemmLaunchSm90 {
using Config = BackwardGemmConfigSm90<Compute>;
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<std::size_t>(Config::kWaveRows) *
static_cast<std::size_t>(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 <int Compute = 90>
struct BackwardTpParamsSm90 {
static constexpr int kCompute = Compute;
BackwardGemmParamsSm90<Compute> 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 <bool ReturnEntropy, int Compute = 90>
void fused_linear_scaled_cross_entropy_backward(
const BackwardTpParamsSm90<Compute>& params,
cudaStream_t stream);
extern template void fused_linear_scaled_cross_entropy_backward<false, 90>(
const BackwardTpParamsSm90<90>&,
cudaStream_t);
extern template void fused_linear_scaled_cross_entropy_backward<true, 90>(
const BackwardTpParamsSm90<90>&,
cudaStream_t);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cu
+753
View File
@@ -0,0 +1,753 @@
#include "backward_reduce_sm90.cuh"
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <device/nvshmemx_collective_launch_apis.h>
#include <cstddef>
#include <cstdint>
#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 <class Layout, class Element>
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 <class Layout, class Element>
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 <class Layout, class Element>
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 <class Layout, class Element>
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 CommConfig, int Compute>
__device__ void direct_peer_communicate(
BackwardDxDwSmemSm90<Compute>& smem,
const BackwardGemmParamsSm90<Compute>& params,
const DxReduceWorkspace<float>& 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<Compute>;
using Config = typename Traits::Config;
int comm_warp =
static_cast<int>(threadIdx.x) / Traits::kWarpSize -
CommConfig::kFirstCommWarp;
int lane = static_cast<int>(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<CommConfig>(unit, num_n_tiles, group_id);
DxCtaGroupSlot slot =
dx_cta_group_slot<CommConfig>(next_index);
std::uint64_t ready_epoch =
static_cast<std::uint64_t>(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<std::size_t>(group.num_tiles) * Traits::kTileM;
std::size_t first_row =
message_rows * static_cast<std::size_t>(comm_warp) /
kDxCommWarpsPerChannel;
std::size_t last_row =
message_rows * static_cast<std::size_t>(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<CommConfig>(cta, slot.comm_warp, slot.stage);
std::size_t ready_offset = dx_sync_offset<CommConfig>(
cta,
comm_warp,
slot.stage,
kDxReadyPhase,
mapping.size);
std::size_t complete_offset = dx_sync_offset<CommConfig>(
cta,
comm_warp,
slot.stage,
kDxCompletePhase,
mapping.size);
std::uint64_t epoch =
dx_epoch_base(comm) |
(static_cast<std::uint64_t>(wave + 1) << 16) |
static_cast<std::uint64_t>(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<int>(
message_vector / kVectorsPerTile);
int vector_in_tile = static_cast<int>(
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<const float4*>(
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<std::size_t>(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<Compute>::kNumThreads, 1)
void backward_dx_dw_direct_peer_wave_kernel_sm90(
__grid_constant__ const Bundle tma,
__grid_constant__ const BackwardGemmParamsSm90<Compute> params,
__grid_constant__ const DxReduceWorkspace<float> comm,
__grid_constant__ const liger_cute::detail::DirectPeerReduceView mapping,
int wave) {
using Traits = BackwardGemmTraitsSm90<Compute>;
using Config = typename Traits::Config;
using Launch = BackwardGemmLaunchSm90<Compute>;
using Smem = BackwardDxDwSmemSm90<Compute>;
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<Smem*>(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<Compute>();
int num_ctas = static_cast<int>(gridDim.x);
int cta = static_cast<int>(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<CommConfig>(dx_n_tiles);
DxPhaseSm90<CommConfig, Compute>::
template produce<ReturnEntropy>(
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<Compute>();
}
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<Compute>::template produce<ReturnEntropy>(
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<CommConfig>(dx_n_tiles);
direct_peer_communicate<CommConfig, Compute>(
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<CommConfig>(dx_n_tiles);
int next_index = 0;
DxPhaseSm90<CommConfig, Compute>::
template consume<ReturnEntropy>(
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<Compute>();
}
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<Compute>::template consume<ReturnEntropy>(
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<Compute>();
}
template <int Compute>
__global__ void cluster_local_allgather_scatter_kernel(
__grid_constant__ const BackwardGemmParamsSm90<Compute> params,
__grid_constant__ const DxReduceWorkspace<float> 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<Compute>;
using Config = typename Traits::Config;
using CommConfig = DxCommConfig<Config, kDxRingStages, 1, Compute>;
int unit = static_cast<int>(blockIdx.x);
if (unit >= tiles_per_wave) return;
int comm_warp = static_cast<int>(threadIdx.x) / Traits::kWarpSize;
int lane = static_cast<int>(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<std::size_t>(wave) * tiles_per_wave + unit) *
packed_tile_elements +
static_cast<std::size_t>(comm_warp) * packed_segment_elements;
int cta = unit % grid_ctas;
int next_index = unit / grid_ctas;
DxCtaGroupSlot slot = dx_cta_group_slot<CommConfig>(next_index);
const float* values = packed_source;
if (mapping.size > 1) {
std::size_t ready_offset = dx_sync_offset<CommConfig>(
cta,
comm_warp,
slot.stage,
kDxReadyPhase,
mapping.size);
std::size_t complete_offset = dx_sync_offset<CommConfig>(
cta,
comm_warp,
slot.stage,
kDxCompletePhase,
mapping.size);
std::uint64_t epoch =
dx_epoch_base(comm) | 0x20000000u |
(static_cast<std::uint64_t>(wave) << 16) |
static_cast<std::uint64_t>(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<std::size_t>(wave) * tiles_per_wave + unit) *
kTileElements;
std::size_t segment_begin =
static_cast<std::size_t>(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<std::size_t>(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<int>(
tile_vector / (Traits::kTileN / 4));
int column = static_cast<int>(
(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<const float4*>(values)[vector];
typename DxPhaseSm90<CommConfig, Compute>::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<std::size_t>(output_row) *
params.hidden +
output_column);
*output = packed;
}
}
}
template <bool RequiresRemote>
void finalize_cluster_typed(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& comm,
const liger_cute::detail::NvlsReduceView& mapping,
const liger_cute::detail::RemoteReduceView& remote,
int grid,
int num_waves,
cudaStream_t stream) {
using CommConfig = DxCommConfig<HostConfig, kDxRingStages, 1, 90>;
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<std::size_t>(num_waves) *
static_cast<std::size_t>(tiles_per_wave) *
CommConfig::kTileElements;
std::size_t packed_elements =
full_elements / static_cast<std::size_t>(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<BackwardGemmParamsSm90<90>*>(&params),
const_cast<DxReduceWorkspace<float>*>(&comm),
const_cast<liger_cute::detail::NvlsReduceView*>(&mapping),
&grid,
&wave,
&tiles_per_wave,
&num_n_tiles};
int status = nvshmemx_collective_launch(
reinterpret_cast<const void*>(kernel),
dim3(static_cast<unsigned>(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<float>& comm,
const liger_cute::detail::DirectPeerReduceView& mapping,
int grid,
cudaStream_t stream,
int wave) {
using CommConfig =
DxCommConfig<HostConfig, kDxRingStages, TilesPerReduce, 90>;
using Smem = BackwardDxDwSmemSm90<90>;
auto* w = static_cast<const HostElement*>(params.weight);
auto* dz = static_cast<HostElement*>(params.dz_workspace);
auto* dw = static_cast<HostElement*>(params.grad_weight);
int padded_vocab = HostLaunch::padded_vocab(params.local_vocab);
auto tma_dz_load =
make_row_major_load<typename HostTraits::SmemAKMajor::Single>(
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<const HostElement*>(params.x),
params.hidden,
params.tokens,
params.hidden);
auto tma_dw_store =
make_row_major_store<typename HostTraits::DwSmemStore1::Single>(
dw, params.local_vocab, params.hidden);
auto tma_dw_add =
make_row_major_reduce_add<typename HostTraits::DwSmemStore1::Single>(
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<int>(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<BackwardGemmParamsSm90<90>*>(&params),
const_cast<DxReduceWorkspace<float>*>(&comm),
const_cast<liger_cute::detail::DirectPeerReduceView*>(&mapping),
&wave};
int status = nvshmemx_collective_launch(
reinterpret_cast<const void*>(kernel),
dim3(static_cast<unsigned>(grid), 1, 1),
dim3(HostConfig::kNumThreads, 1, 1),
args,
static_cast<std::size_t>(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<float>& 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 <bool RequiresRemote>
void finalize_backward_dx_cluster_sm90(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& comm,
const liger_cute::detail::NvlsReduceView& mapping,
const liger_cute::detail::RemoteReduceView& remote,
int grid,
int num_waves,
cudaStream_t stream) {
finalize_cluster_typed<RequiresRemote>(
params, comm, mapping, remote, grid, num_waves, stream);
}
template void finalize_backward_dx_cluster_sm90<false>(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& 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<true>(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& 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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_reduce_sm90.cuh
+51
View File
@@ -0,0 +1,51 @@
#pragma once
// FSLCE-specific adapters around the reusable tensor-parallel reduction API.
#include <cuda_runtime.h>
#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<float>& comm,
const liger_cute::detail::DirectPeerReduceView& mapping,
int grid,
cudaStream_t stream,
int wave);
template <bool RequiresRemote>
void finalize_backward_dx_cluster_sm90(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& 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<false>(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& 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<true>(
const BackwardGemmParamsSm90<90>& params,
const DxReduceWorkspace<float>& 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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_sm90_selftest.cu
+701
View File
@@ -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=<n> 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 <arpa/inet.h>
#include <nvshmem.h>
#include <nvshmemx.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#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<const char*>(data);
while (bytes != 0) {
ssize_t sent = send(fd, cursor, bytes, 0);
if (sent <= 0) std::exit(1);
cursor += sent;
bytes -= static_cast<std::size_t>(sent);
}
}
void recv_all(int fd, void* data, std::size_t bytes) {
char* cursor = static_cast<char*>(data);
while (bytes != 0) {
ssize_t received = recv(fd, cursor, bytes, 0);
if (received <= 0) std::exit(1);
cursor += received;
bytes -= static_cast<std::size_t>(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<std::uint16_t>(port));
if (bind(server, reinterpret_cast<sockaddr*>(&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<std::uint16_t>(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<sockaddr*>(&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<std::uint32_t>(state >> 33);
return (static_cast<float>(bits % 2048) / 1024.0f - 1.0f) * scale;
}
struct Reference {
std::vector<float> dx; // [tokens, hidden], TP-summed
std::vector<float> dw_local; // [local_vocab, hidden], this shard
std::vector<float> lse; // [tokens]
std::vector<float> 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<float>& x,
const std::vector<float>& weight_global,
const std::vector<std::int64_t>& target,
const std::vector<float>& grad_output,
const std::vector<float>& 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<std::size_t>(tokens) * hidden, 0.0f);
out.dw_local.assign(
static_cast<std::size_t>(local_vocab) * hidden, 0.0f);
out.lse.assign(tokens, 0.0f);
out.entropy.assign(tokens, 0.0f);
std::vector<float> scaled(global_vocab);
std::vector<float> 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<std::size_t>(m) * hidden + k] *
weight_global[static_cast<std::size_t>(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<double>(scaled[v]) - static_cast<double>(row_max));
sum += e;
weighted += e * static_cast<double>(scaled[v]);
}
float lse = row_max + static_cast<float>(std::log(sum));
out.lse[m] = lse;
out.entropy[m] = lse - static_cast<float>(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<int>(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<std::size_t>(global_v) * hidden + h];
}
reduced += partial;
}
out.dx[static_cast<std::size_t>(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<std::size_t>(v) * hidden + h] +=
value * x[static_cast<std::size_t>(m) * hidden + h];
}
}
}
return out;
}
struct Comparison {
int bad;
float max_absolute;
float max_relative;
};
Comparison compare(
const char* label,
const std::vector<float>& actual,
const std::vector<float>& 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<TestCommConfig>(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 <bool ReturnEntropy>
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<float>(hidden));
std::vector<float> host_x(static_cast<std::size_t>(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<float> host_w(
static_cast<std::size_t>(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<std::int64_t> host_target(tokens);
std::vector<float> host_grad_output(tokens);
std::vector<float> 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<std::int64_t>(
(state >> 17) % static_cast<std::uint64_t>(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<Element> 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<Element> device_w_host(
static_cast<std::size_t>(local_vocab) * hidden);
for (std::size_t i = 0; i < device_w_host.size(); ++i) {
device_w_host[i] = __float2bfloat16(
host_w[static_cast<std::size_t>(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<float> 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<Element> raw_dx(device_x_host.size());
std::vector<Element> 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<float> got_dx(raw_dx.size());
std::vector<float> 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<int>(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<std::int64_t>(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<false>("dX/dW aligned, no entropy", small, 2, 2);
run_case<true>("dX/dW aligned, entropy", small, 2, 2);
run_case<true>("dX/dW ragged M/H/V, TilesPerReduce=1", ragged, 1, 2);
run_case<true>("dX/dW ragged M/H/V, TilesPerReduce=2", ragged, 2, 2);
run_case<true>("dX/dW ragged M/H/V, TilesPerReduce=4", ragged, 4, 2);
run_case<true>("dX/dW ragged, legacy channel arg=1", ragged, 2, 1);
run_case<true>("dX/dW ragged, legacy channel arg=4", ragged, 2, 4);
run_case<false>("dX/dW two waves, TilesPerReduce=2", two_waves, 2, 4);
run_case<true>(
"dX/dW H2048 odd K tiles and temperature",
odd_k_tiles,
1,
2);
run_case<false>("dX/dW full resident CTA grid", full_grid, 1, 2);
run_case<false>(
"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<int*>(
pool.get_symmetric("fslce_bwd_selftest_verdict_local", sizeof(int)));
int* total_slot = static_cast<int*>(
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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/config.cuh
+62
View File
@@ -0,0 +1,62 @@
#pragma once
#include <cuda_runtime.h>
#include <cstdint>
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 <typename Config>
__host__ __device__ constexpr int dx_groups_per_m_tile(int num_n_tiles) {
return ceil_div(num_n_tiles, Config::kCoalesceTiles);
}
template <typename Config>
__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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/dx_reduce.cuh
+240
View File
@@ -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 <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
#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 <typename Element>
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 <typename Element>
__device__ __forceinline__ std::uint64_t dx_epoch_base(
const DxReduceWorkspace<Element>& workspace) {
return *workspace.launch_epoch;
}
#endif
template <typename Config>
__host__ __device__ inline std::size_t dx_slot_index(
int cta, int comm_warp, int stage) {
return (static_cast<std::size_t>(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 <typename Config>
__host__ __device__ inline std::size_t dx_slot_offset(
int cta, int comm_warp, int stage) {
std::size_t slot = dx_slot_index<Config>(cta, comm_warp, stage);
return slot * static_cast<std::size_t>(Config::kTilesPerReduce) *
static_cast<std::size_t>(Config::kTileElements);
}
template <typename Config>
__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<Config>(cta, comm_warp, stage);
return (slot * kDxSyncPhases + static_cast<std::size_t>(phase)) *
static_cast<std::size_t>(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 <typename Config>
__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 <typename Config>
__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<Config>(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<Config>(
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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_mainloop_sm90.cuh
+794
View File
@@ -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 <cstdint>
namespace liger {
namespace fused_scaled_linear_cross_entropy {
using namespace cute;
// ───────────────────────────────────────────────────────────────────────────
// Traits
// ───────────────────────────────────────────────────────────────────────────
template <int Compute = 90>
struct ForwardGemmTraitsSm90 {
static_assert(Compute == 90,
"SM90 fused scaled linear cross entropy requires Compute=90");
using Config = ForwardGemmConfigSm90<Compute>;
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<Int<kClusterM>, _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<GMMA::Major::K, GMMA::Major::K>>,
"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<Element, kTileM, kTileK, kStages, Compute>;
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<kStages, Compute>;
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 <int Compute, bool ReturnEntropy>
struct ForwardGemmSmemSm90 {
using Traits = ForwardGemmTraitsSm90<Compute>;
using Element = typename Traits::Element;
using ElementLogit = typename Traits::ElementLogit;
static constexpr int kSmemX = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int kSmemW = cosize_v<typename Traits::SmemLayoutW>;
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 <int Compute = 90>
CUTE_DEVICE void forward_epilogue_barrier_sm90() {
using Config = ForwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
CUTE_DEVICE ForwardGemmWorkSm90 forward_gemm_assign_work_sm90(
const ForwardGemmSplitSm90<Compute>& split,
int cluster_work,
int cluster_rank) {
using Config = ForwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct ForwardGemmProducerSm90 {
using Traits = ForwardGemmTraitsSm90<Compute>;
static constexpr int kCompute = Compute;
static_assert(Compute == 90,
"SM90 fused scaled linear cross entropy requires Compute=90");
template <bool ReturnEntropy, class Pipeline, class TmaLoadX, class TmaLoadW>
CUTE_DEVICE static void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
ForwardGemmSmemSm90<Compute, ReturnEntropy>& smem,
TmaLoadX const& tma_load_x,
TmaLoadW const& tma_load_w,
ForwardGemmParamsSm90<Compute> 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<int64_t>(params.tokens),
static_cast<int64_t>(params.hidden)));
auto mW = tma_load_w.get_tma_tensor(make_shape(
static_cast<int64_t>(params.local_vocab),
static_cast<int64_t>(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<Traits::kTileM>{}, Int<Traits::kTileK>{}),
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<Traits::kAccumulatorTileN>{}, Int<Traits::kTileK>{}),
make_coord(fragment_base + 0, _));
auto gW1 = local_tile(
mW,
make_tile(
Int<Traits::kAccumulatorTileN>{}, Int<Traits::kTileK>{}),
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 <int Compute = 90>
struct ForwardGemmConsumerSm90 {
using Traits = ForwardGemmTraitsSm90<Compute>;
using Config = typename Traits::Config;
using Epilogue = ForwardGemmEpilogueSm90<Compute>;
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 <bool ReturnEntropy, int SliceBase, class Accum, class CoordC>
CUTE_DEVICE static void stage_slice(
ForwardGemmSmemSm90<Compute, ReturnEntropy>& 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<typename Traits::ElementLogit>(accum(i));
}
}
cutlass::arch::fence_view_async_shared();
asm volatile("fence.acq_rel.cta;" ::: "memory");
forward_epilogue_barrier_sm90<Compute>();
}
// _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 <bool ReturnEntropy>
CUTE_DEVICE static void fold_slice(
ForwardGemmSmemSm90<Compute, ReturnEntropy>& 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<float>(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<float>(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<ReturnEntropy, kChunkN>(
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<float>(logits[target_col - slice_vocab_base]) *
inverse_temperature;
softmax.has_target = 1;
}
}
template <bool ReturnEntropy, class Pipeline>
CUTE_DEVICE static void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
ForwardGemmSmemSm90<Compute, ReturnEntropy>& smem,
ForwardGemmParamsSm90<Compute> const& params,
ForwardGemmPartialsSm90<Compute> const& partials,
ForwardGemmWorkSm90 const& work,
int num_split_tiles,
int num_k_tiles) {
typename Traits::TiledMma tiled_mma;
int tid_in_mma =
static_cast<int>(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<Traits::kTileM>, Int<Traits::kAccumulatorTileN>>;
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<std::int64_t>(params.local_vocab)) {
target_col = static_cast<int>(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<ReturnEntropy, 0>(smem, accum0, coord_c, 0);
stage_slice<ReturnEntropy, kSliceN>(smem, accum0, coord_c, 1);
fold_slice<ReturnEntropy>(
smem, softmax, 0, epi_row, epi_segment,
logical_base,
local_vocab - logical_base,
target_col, params.inverse_temperature);
forward_epilogue_barrier_sm90<Compute>();
stage_slice<ReturnEntropy, 0>(smem, accum1, coord_c, 0);
fold_slice<ReturnEntropy>(
smem, softmax, 1, epi_row, epi_segment,
logical_base + kSliceN,
local_vocab - (logical_base + kSliceN),
target_col, params.inverse_temperature);
forward_epilogue_barrier_sm90<Compute>();
stage_slice<ReturnEntropy, kSliceN>(smem, accum1, coord_c, 1);
fold_slice<ReturnEntropy>(
smem, softmax, 0, epi_row, epi_segment,
logical_base + kAccumN,
local_vocab - (logical_base + kAccumN),
target_col, params.inverse_temperature);
forward_epilogue_barrier_sm90<Compute>();
fold_slice<ReturnEntropy>(
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<Compute>();
}
// 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<Compute>();
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<ReturnEntropy>(
softmax, other);
Epilogue::template store_partial<ReturnEntropy>(
combined,
partials,
work.output_work * Traits::kTileM + epi_row);
}
}
};
// ───────────────────────────────────────────────────────────────────────────
// Fragment kernel
// ───────────────────────────────────────────────────────────────────────────
template <bool ReturnEntropy, int Compute, class TmaLoadX, class TmaLoadW>
__global__ __launch_bounds__(ForwardGemmTraitsSm90<Compute>::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<Compute> params,
__grid_constant__ const ForwardGemmPartialsSm90<Compute> partials,
__grid_constant__ const ForwardGemmSplitSm90<Compute> split) {
using Traits = ForwardGemmTraitsSm90<Compute>;
using Config = typename Traits::Config;
using Launch = ForwardGemmLaunchSm90<Compute>;
using Smem = ForwardGemmSmemSm90<Compute, ReturnEntropy>;
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<Smem*>(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<Compute>();
// 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<Compute>(
split,
static_cast<int>(blockIdx.z),
static_cast<int>(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<Compute>::template run<ReturnEntropy>(
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<Compute>(
split,
static_cast<int>(blockIdx.z),
static_cast<int>(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<Compute>::template run<ReturnEntropy>(
pipe,
consumer_state,
smem,
params,
partials,
work,
num_split_tiles,
num_k_tiles);
}
sm90::cluster_exit_sm90<Compute>();
}
// ───────────────────────────────────────────────────────────────────────────
// Split reducer — the reference finalizer, stopped at *local* statistics so
// The host forward wrapper reduces across the tensor-parallel group.
// ───────────────────────────────────────────────────────────────────────────
template <bool ReturnEntropy, int Compute>
__global__ __launch_bounds__(ForwardGemmConfigSm90<Compute>::kTileM, 1)
void forward_split_reduce_kernel_sm90(
__grid_constant__ const ForwardGemmParamsSm90<Compute> params,
__grid_constant__ const ForwardGemmPartialsSm90<Compute> partials,
__grid_constant__ const ForwardGemmSplitSm90<Compute> split) {
using Config = ForwardGemmConfigSm90<Compute>;
using Epilogue = ForwardGemmEpilogueSm90<Compute>;
int tid = static_cast<int>(threadIdx.x);
int pid_m = static_cast<int>(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<std::int64_t>(params.local_vocab)) ? 1 : 0;
Epilogue::template store_row<ReturnEntropy>(state, params.output, row);
}
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90.cuh
+523
View File
@@ -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 <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
#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 <int Compute = 90>
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 <int Compute = 90>
struct ForwardGemmTuningSm90 {
using Config = ForwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
struct ForwardGemmSplitSm90 {
using Config = ForwardGemmConfigSm90<Compute>;
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 <int Compute = 90>
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 <int Compute = 90>
struct ForwardGemmParamsSm90 {
using Config = ForwardGemmConfigSm90<Compute>;
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<Compute>::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<Compute> tuning = {};
};
// ───────────────────────────────────────────────────────────────────────────
// Online-softmax epilogue contract
// ───────────────────────────────────────────────────────────────────────────
template <int Compute = 90>
struct ForwardGemmEpilogueSm90 {
using Config = ForwardGemmConfigSm90<Compute>;
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 <bool ReturnEntropy, int FragmentSize>
__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<ReturnEntropy>(
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 <bool ReturnEntropy, int ChunkN>
__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 <bool ReturnEntropy>
__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 <bool ReturnEntropy>
__host__ __device__ static void store_partial(
const OnlineSoftmaxState& state,
const ForwardGemmPartialsSm90<Compute>& 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 <bool ReturnEntropy>
__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 <int Compute = 90>
struct ForwardGemmLaunchSm90 {
using Config = ForwardGemmConfigSm90<Compute>;
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<Compute> resolve_split(
const ForwardGemmTuningSm90<Compute>& tuning,
int tokens,
int local_vocab,
int max_active_cluster_pairs) {
ForwardGemmSplitSm90<Compute> 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<std::size_t>(num_m_tiles(tokens)) *
static_cast<std::size_t>(cap) *
static_cast<std::size_t>(Config::kTileM);
return rows * sizeof(float) * (return_entropy ? 4u : 3u);
}
// Mirrors ForwardGemmSmemSm90<Compute, ReturnEntropy> — 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 <int Compute = 90>
struct ForwardTpParamsSm90 {
static constexpr int kCompute = Compute;
ForwardGemmParamsSm90<Compute> 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 <bool ReturnEntropy, int Compute = 90>
void fused_linear_scaled_cross_entropy_forward(
const ForwardTpParamsSm90<Compute>& params,
cudaStream_t stream);
extern template void fused_linear_scaled_cross_entropy_forward<false, 90>(
const ForwardTpParamsSm90<90>&,
cudaStream_t);
extern template void fused_linear_scaled_cross_entropy_forward<true, 90>(
const ForwardTpParamsSm90<90>&,
cudaStream_t);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_sm90_selftest.cu
+414
View File
@@ -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<this dir> -I<core include> -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 <cuda_runtime.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
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<std::uint32_t>(state >> 33);
return static_cast<float>(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<ReferenceRow> reference(
const std::vector<Element>& x,
const std::vector<Element>& weight,
const std::vector<std::int64_t>& target,
int tokens,
int hidden,
int local_vocab,
std::int64_t vocab_start,
std::int64_t ignore_index,
float inverse_temperature) {
std::vector<ReferenceRow> rows(tokens);
std::vector<float> 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<float>(x[m * hidden + k]) *
static_cast<float>(weight[v * hidden + k]);
}
float staged = static_cast<float>(static_cast<ElementLogit>(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<double>(scaled[v]) - static_cast<double>(row_max));
sum += e;
weighted += e * static_cast<double>(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<int>(local)];
}
}
rows[m].max_value = row_max;
rows[m].exp_sum = static_cast<float>(sum);
rows[m].lse = row_max + static_cast<float>(std::log(sum));
rows[m].target_logit = target_logit;
rows[m].weighted_sum = static_cast<float>(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 <bool ReturnEntropy>
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<const Element*>(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<const Element*>(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<int>(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<std::size_t>(split.num_m_tiles) *
static_cast<std::size_t>(split.split_n) *
static_cast<std::size_t>(Config::kTileM);
float* scratch = static_cast<float*>(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<unsigned>(Config::kClusterM),
1u,
static_cast<unsigned>(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<ReturnEntropy, 90>
<<<split.num_m_tiles, Config::kTileM, 0, stream>>>(
params, partials, split);
}
template <bool ReturnEntropy>
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<Element> host_x(static_cast<std::size_t>(tokens) * hidden);
std::vector<Element> host_w(static_cast<std::size_t>(local_vocab) * hidden);
std::vector<std::int64_t> host_target(tokens);
const float scale = 1.0f / std::sqrt(static_cast<float>(hidden));
for (auto& value : host_x) {
value = static_cast<Element>(next_random(state) * scale * 4.0f);
}
for (auto& value : host_w) {
value = static_cast<Element>(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<std::int64_t>(bits % static_cast<std::uint64_t>(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<ReturnEntropy>(params, nullptr);
CUDA_OK(cudaGetLastError());
CUDA_OK(cudaDeviceSynchronize());
std::vector<float> host_stats(4 * tokens);
CUDA_OK(cudaMemcpy(host_stats.data(), device_stats,
host_stats.size() * sizeof(float), cudaMemcpyDeviceToHost));
std::vector<ReferenceRow> 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<long long>(vocab_start),
temperature,
static_cast<int>(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<false>("aligned / no entropy", 256, 128, 640, 0, 1.0f, 1);
run_case<true>("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<false>("ragged M/H/V / no entropy", 200, 104, 517, 0, 1.0f, 3);
run_case<true>("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<true>("ragged K tail", 128, 72, 333, 0, 1.0f, 8);
// Sharded vocabulary with a non-zero start plus temperature scaling.
run_case<true>("shard offset + temperature", 384, 192, 901, 4096, 0.7f, 5);
run_case<false>("single M tile", 96, 64, 321, 12288, 1.3f, 6);
// Local vocabulary smaller than one logical N320 tile.
run_case<true>("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<true>("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<true>("split: uneven base=2 extra=3",
1024, 128, 2560, 0, 1.0f, 10, uneven);
run_case<false>("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<true>("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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cu
+132
View File
@@ -0,0 +1,132 @@
#include "forward_reduce.cuh"
#include <cstddef>
#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<std::size_t>(tokens) * sizeof(float);
}
std::size_t packed_bytes_at(int tokens) {
return static_cast<std::size_t>(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 <bool ReturnEntropy>
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<float*>(
pool.get_device(Names::kLocalMax, token_bytes));
workspace.global_max = static_cast<float*>(
pool.get_device(Names::kGlobalMax, token_bytes));
workspace.packed = static_cast<float*>(
pool.get_device(Names::kPacked, packed_bytes));
workspace.reduced = static_cast<float*>(
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<float*>(
pool.get_device(Names::kLocalSum, token_bytes));
workspace.local.local_target = static_cast<float*>(
pool.get_device(Names::kLocalTarget, token_bytes));
workspace.local.local_weighted_sum = static_cast<float*>(
pool.get_device(Names::kLocalWeighted, token_bytes));
workspace.tokens = tokens;
workspace.fields = forward_reduced_fields(ReturnEntropy);
return workspace;
}
template ForwardTpWorkspace reserve_forward_tp_workspace<false>(int);
template ForwardTpWorkspace reserve_forward_tp_workspace<true>(int);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_reduce.cuh
+80
View File
@@ -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 <cuda_runtime.h>
#include <cstddef>
#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 <bool ReturnEntropy>
ForwardTpWorkspace reserve_forward_tp_workspace(int tokens);
extern template ForwardTpWorkspace
reserve_forward_tp_workspace<false>(int);
extern template ForwardTpWorkspace
reserve_forward_tp_workspace<true>(int);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_tp_sm90_selftest.cu
+427
View File
@@ -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 <includes>
// nvcc -std=c++17 -arch=sm_90a -O3 --expt-relaxed-constexpr -rdc=true \
// -c forward_reduce.cu -o forward_reduce.o <includes>
// 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<this dir> -I<core include> -I<core src/moe> \
// -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 <nvshmem.h>
#include <nvshmemx.h>
#include <nccl.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstring>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
#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<std::uint32_t>(state >> 33);
return static_cast<float>(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<std::uint32_t>(state >> 33);
return (static_cast<float>(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<ReferenceRow> reference(
const std::vector<float>& x,
const std::vector<float>& weight_global,
const std::vector<std::int64_t>& target,
int tokens,
int hidden,
int global_vocab,
std::int64_t ignore_index,
float inverse_temperature) {
std::vector<ReferenceRow> rows(tokens);
std::vector<float> 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<std::size_t>(m) * hidden + k] *
weight_global[static_cast<std::size_t>(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<double>(scaled[v]) - static_cast<double>(row_max));
sum += e;
weighted += e * static_cast<double>(scaled[v]);
}
float lse = row_max + static_cast<float>(std::log(sum));
bool ignored = target[m] == ignore_index;
rows[m].expectation = static_cast<float>(weighted / sum);
rows[m].row_max = row_max;
rows[m].sum = static_cast<float>(sum);
rows[m].lse = lse;
rows[m].nll = ignored ? 0.0f : lse - scaled[target[m]];
rows[m].entropy = ignored
? 0.0f
: lse - static_cast<float>(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 <bool ReturnEntropy>
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<std::int64_t>(g_pe) * local_vocab;
const float scale = 4.0f / std::sqrt(static_cast<float>(hidden));
// Rebuild the global tensors identically on every PE.
std::vector<float> host_x(static_cast<std::size_t>(tokens) * hidden);
for (std::size_t i = 0; i < host_x.size(); ++i) {
host_x[i] = element_at(seed, i, scale);
}
std::vector<float> host_w(
static_cast<std::size_t>(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<std::int64_t> 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<std::int64_t>(
(target_state >> 17) % static_cast<std::uint64_t>(global_vocab));
}
}
// Device operands: X is replicated, W is this PE's contiguous shard.
std::vector<Element> 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<Element> device_w_host(
static_cast<std::size_t>(local_vocab) * hidden);
for (std::size_t i = 0; i < device_w_host.size(); ++i) {
device_w_host[i] = __float2bfloat16(
host_w[static_cast<std::size_t>(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<std::int64_t>(g_nccl_comm);
fslce::fused_linear_scaled_cross_entropy_forward<
ReturnEntropy, 90>(params, nullptr);
CUDA_OK(cudaGetLastError());
CUDA_OK(cudaDeviceSynchronize());
std::vector<float> 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<ReturnEntropy>(tokens);
std::vector<float> host_global_max(tokens);
std::vector<float> host_reduced(
static_cast<std::size_t>(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<ReferenceRow> 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<int>(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<false>("tp aligned / no entropy", 256, 128, 640, 1.0f, 101);
run_case<true>("tp aligned / entropy", 256, 128, 640, 1.0f, 102);
run_case<true>("tp ragged M/H/V", 200, 104, 517, 0.8f, 103);
run_case<false>("tp temperature / no entropy", 384, 192, 901, 1.4f, 104);
run_case<true>("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<int*>(
pool.get_symmetric("fslce_tp_selftest_verdict_local", sizeof(int)));
int* total_slot = static_cast<int*>(
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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward.cu
+542
View File
@@ -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 <cstddef>
#include <cstdint>
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 <class Layout, class Element>
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 <class Layout, int TileRows, int TileCols, class Element>
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<TileRows>{}, Int<TileCols>{}),
_1{});
}
template <class Layout, class Element>
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 <class Layout, class Element>
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 <class Layout, int TileRows, int TileCols, class Element>
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<TileRows>{}, Int<TileCols>{}),
_1{});
}
template <class Layout, class Element>
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<TileRows>{}, Int<TileCols>{}),
Int<Multicast>{});
}
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<std::uintptr_t>(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<std::uintptr_t>(gemm.grad_input) % 8 == 0,
"grad_input must be 8 B aligned for vectorized BF16 stores");
LIGER_CHECK(
(static_cast<std::size_t>(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<HostConfig, kDxRingStages, TilesPerReduce, 90>;
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<const HostElement*>(gemm.x);
auto* w = static_cast<const HostElement*>(gemm.weight);
auto* dz = static_cast<HostElement*>(gemm.dz_workspace);
auto* dw_out = static_cast<HostElement*>(gemm.grad_weight);
DxReduceWorkspace<float> 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<float*>(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<typename HostDxClusterTraits::SmemA::Single>(
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<int>(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<int>(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<std::size_t>(num_waves) *
HostConfig::kWaveRows *
static_cast<std::size_t>(
HostLaunch::num_dx_n_tiles(gemm.hidden)) *
HostDxClusterTraits::kTileN /
static_cast<std::size_t>(
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<unsigned>(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<unsigned>(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<RequiresRemote>(
device_params,
comm,
reduce.nvls,
reduce.remote,
cluster_grid,
num_waves,
stream);
}
liger_cute::detail::end_tp_reduce(stream);
}
template <int TilesPerReduce, bool ReturnEntropy>
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 <bool ReturnEntropy, int Compute>
void fused_linear_scaled_cross_entropy_backward(
const BackwardTpParamsSm90<Compute>& 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<false, 90>(
const BackwardTpParamsSm90<90>&,
cudaStream_t);
template void fused_linear_scaled_cross_entropy_backward<true, 90>(
const BackwardTpParamsSm90<90>&,
cudaStream_t);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward.cu
+348
View File
@@ -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 <nccl.h>
#include <cstddef>
#include <cstdint>
namespace liger {
namespace fused_scaled_linear_cross_entropy {
namespace {
template <bool ReturnEntropy>
__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<int>(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 <bool ReturnEntropy>
__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<int>(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<ReturnEntropy>(
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 <bool ReturnEntropy, int Compute>
void fused_linear_scaled_cross_entropy_forward(
const ForwardTpParamsSm90<Compute>& params,
cudaStream_t stream) {
using Traits = ForwardGemmTraitsSm90<Compute>;
using Config = typename Traits::Config;
using Launch = ForwardGemmLaunchSm90<Compute>;
using Element = typename Traits::Element;
using Smem = ForwardGemmSmemSm90<Compute, ReturnEntropy>;
static_assert(Compute == 90,
"SM90 fused scaled linear cross entropy requires Compute=90");
const ForwardGemmParamsSm90<Compute>& 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<std::uintptr_t>(input.x) % 16 == 0 &&
reinterpret_cast<std::uintptr_t>(input.weight) % 16 == 0,
"forward GEMM operands must be 16 B aligned for TMA");
LIGER_CHECK(
(static_cast<std::size_t>(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<ReturnEntropy>(input.tokens);
ForwardGemmParamsSm90<Compute> 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<Compute>::template make_load<
typename Traits::SmemLayoutX1,
Traits::kTileM,
Traits::kTileK>(
static_cast<const Element*>(gemm_params.x),
gemm_params.tokens,
gemm_params.hidden);
auto tma_load_w =
sm90::TmaSm90<Compute>::template make_load_multicast<
typename Traits::SmemLayoutW1,
Traits::kAccumulatorTileN,
Traits::kTileK,
Traits::kClusterM>(
static_cast<const Element*>(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<int>(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<Compute>;
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<Compute> 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<std::size_t>(split.num_m_tiles) *
static_cast<std::size_t>(split.split_n) *
static_cast<std::size_t>(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<std::uintptr_t>(gemm_params.workspace) % 16 == 0,
"forward GEMM workspace must be 16 B aligned");
float* scratch = static_cast<float*>(gemm_params.workspace);
ForwardGemmPartialsSm90<Compute> 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<unsigned>(Config::kClusterM),
1u,
static_cast<unsigned>(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<ReturnEntropy, Compute>
<<<split.num_m_tiles, Config::kTileM, 0, stream>>>(
gemm_params, partials, split);
check_cuda(
cudaGetLastError(), "forward_split_reduce_kernel_sm90 launch");
ncclComm_t comm =
reinterpret_cast<ncclComm_t>(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<ReturnEntropy>
<<<post_blocks, kPostThreads, 0, stream>>>(
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<std::size_t>(reduce_workspace.fields) *
static_cast<std::size_t>(gemm_params.tokens),
ncclFloat32,
ncclSum,
comm,
stream),
"ncclAllReduce(SUM)");
finalize_forward_kernel<ReturnEntropy>
<<<post_blocks, kPostThreads, 0, stream>>>(
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<false, 90>(
const ForwardTpParamsSm90<90>&,
cudaStream_t);
template void fused_linear_scaled_cross_entropy_forward<true, 90>(
const ForwardTpParamsSm90<90>&,
cudaStream_t);
} // namespace fused_scaled_linear_cross_entropy
} // namespace liger
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_scaled_linear_cross_entropy.cuh
+10
View File
@@ -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"
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/gemm_sm90.cuh
+506
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/arch/cluster_sm90.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>
#include <cutlass/cutlass.h>
#include <cutlass/arch/reg_reconfig.h>
#include <cutlass/numeric_types.h>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cuda_runtime.h>
#include <cstdint>
namespace liger {
namespace sm90 {
using namespace cute;
// ───────────────────────────────────────────────────────────────────────────
// Architecture constants
// ───────────────────────────────────────────────────────────────────────────
template <int Compute = 90>
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<kAtomTileM>, Int<AtomTileN>, Int<kAtomTileK>>,
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<Atom>, Layout<Shape<Int<NumWarpGroups>, _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 <class Element, int Rows, int Cols, int Stages, int Compute = 90>
struct SmemLayoutKMajorSm90 {
static_assert(Compute == 90, "SmemLayoutKMajorSm90 requires Compute=90");
static_assert(Cols * static_cast<int>(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<Element>;
using Single = decltype(tile_to_shape(
Atom{}, Shape<Int<Rows>, Int<Cols>>{}));
using Staged = decltype(tile_to_shape(
Atom{}, Shape<Int<Rows>, Int<Cols>, Int<Stages>>{}));
static constexpr int kStageBytes = static_cast<int>(
cute::size(Single{}) * sizeof(Element));
static constexpr int kStagedElements = cosize_v<Staged>;
};
// ───────────────────────────────────────────────────────────────────────────
// TMA descriptors
// ───────────────────────────────────────────────────────────────────────────
template <int Compute = 90>
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 <class Element>
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<int64_t>(rows), static_cast<int64_t>(cols)),
make_stride(static_cast<int64_t>(cols), _1{})));
}
template <class SmemLayout, int TileRows, int TileCols, class Element>
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<TileRows>{}, Int<TileCols>{}),
_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<TileRows>{}, Int<TileCols>{}),
Int<Multicast>{});
}
};
// Multicast mask over the cluster's M mode: every CTA sharing this N tile.
template <class ClusterShape, int Compute = 90>
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<ClusterShape>{};
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 <int Stages, int Compute = 90>
struct MainloopPipelineSm90 {
static_assert(Compute == 90, "MainloopPipelineSm90 requires Compute=90");
static constexpr int kCompute = Compute;
static constexpr int kStages = Stages;
using Type = cutlass::PipelineTmaAsync<Stages>;
using State = cutlass::PipelineState<Stages>;
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 <class ClusterShape>
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 <class ClusterShape>
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 <class ClusterShape>
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<Type>();
}
};
// ───────────────────────────────────────────────────────────────────────────
// Warp-group register reconfiguration and cluster synchronisation
// ───────────────────────────────────────────────────────────────────────────
template <int ProducerRegisters, int ConsumerRegisters, int Compute = 90>
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<ProducerRegisters>();
}
CUTE_DEVICE static void consumer() {
cutlass::arch::warpgroup_reg_alloc<ConsumerRegisters>();
}
};
// pipeline_init_arrive / pipeline_init_wait: the cluster-wide handshake that
// must complete before any mbarrier in the cluster is touched.
template <int Compute = 90>
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 <int Compute = 90>
CUTE_DEVICE void cluster_exit_sm90() {
static_assert(Compute == 90, "cluster_exit_sm90 requires Compute=90");
cute::cluster_sync();
}
template <int BarrierId, int NumThreads, int Compute = 90>
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 <int Compute = 90>
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 <class Pipeline, class State, class IssueCopies>
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 <class Pipeline, class State, class IssueMma, class FenceAccumulators>
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 <int Compute = 90>
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 <class Kernel>
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<unsigned>(cluster_m);
attribute.val.clusterDim.y = 1;
attribute.val.clusterDim.z = 1;
cudaLaunchConfig_t config = {};
config.gridDim = grid;
config.blockDim = dim3(static_cast<unsigned>(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 <class Kernel>
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<unsigned>(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 <class Kernel, class... Args>
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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/online_softmax.cuh
+118
View File
@@ -0,0 +1,118 @@
#pragma once
#include <cuda_runtime.h>
#include <cmath>
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 <bool ReturnEntropy>
__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 <bool ReturnEntropy>
__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 <bool ReturnEntropy>
__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 <bool ReturnEntropy>
__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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/state.h
+11
View File
@@ -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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/tma_store_sm90.cuh
+36
View File
@@ -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 <cuda_runtime.h>
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 <int Pending>
__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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cu
+429
View File
@@ -0,0 +1,429 @@
#include "workspace.cuh"
#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
#include <limits>
#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<std::size_t>::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<std::size_t>(Config::kTileM),
static_cast<std::size_t>(Config::kDxTileN),
"dX tile");
std::size_t elements = checked_multiply(
tile_elements,
static_cast<std::size_t>(tiles_per_reduce),
"dX group");
elements = checked_multiply(
elements, static_cast<std::size_t>(kConfiguredStages), "dX ring");
elements = checked_multiply(
elements,
static_cast<std::size_t>(kDxCommWarpsPerChannel),
"dX comm warps");
elements = checked_multiply(
elements, static_cast<std::size_t>(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<std::size_t>(num_ctas),
static_cast<std::size_t>(kDxCommWarpsPerChannel),
"dX sync CTAs");
entries = checked_multiply(
entries, static_cast<std::size_t>(kConfiguredStages), "dX sync ring");
entries = checked_multiply(
entries, static_cast<std::size_t>(kDxSyncPhases), "dX sync phases");
entries = checked_multiply(
entries, static_cast<std::size_t>(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<std::size_t>(tokens) + Config::kWaveRows - 1) /
Config::kWaveRows;
std::size_t rows = checked_multiply(
waves, static_cast<std::size_t>(Config::kWaveRows), "dX durable rows");
std::size_t n_tiles =
(static_cast<std::size_t>(hidden) + Config::kDxTileN - 1) /
Config::kDxTileN;
std::size_t columns = checked_multiply(
n_tiles,
static_cast<std::size_t>(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<std::size_t>(local_partition_size) - 1) /
static_cast<std::size_t>(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<float*>(
pool.get_symmetric(Names::kDxPartial, g_staging_bytes));
auto* reduced = static_cast<float*>(
pool.get_symmetric(Names::kDxReduced, g_reduced_bytes));
auto* reduced_shard = static_cast<float*>(
pool.get_symmetric(
Names::kDxReducedShard, g_packed_durable_bytes));
auto* remote_inbox = static_cast<float*>(
pool.get_symmetric(
Names::kDxRemoteInbox, g_packed_durable_bytes));
auto* remote_signals = static_cast<std::uint64_t*>(
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<std::uint64_t*>(
pool.get_symmetric(Names::kDxSync, sync_bytes));
auto** peer_partial_storage = static_cast<float**>(
pool.get_device(
Names::kDxPeerPartialPointers,
static_cast<std::size_t>(team_size) * sizeof(float*)));
auto** peer_sync_storage = static_cast<std::uint64_t**>(
pool.get_device(
Names::kDxPeerSyncPointers,
static_cast<std::size_t>(team_size) *
sizeof(std::uint64_t*)));
pool.get_device(
Names::kDzWorkspace, Launch::dz_workspace_bytes(max_local_vocab));
auto* launch_epoch = static_cast<std::uint64_t*>(
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<std::size_t>(team_size) *
(2 * sizeof(void*));
return Launch::dz_workspace_bytes(max_local_vocab) +
sizeof(std::uint64_t) + peer_mapping;
}
DxReduceWorkspace<float> 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<float> workspace = {};
workspace.partial = static_cast<float*>(
pool.get_symmetric(Names::kDxPartial, g_staging_bytes));
workspace.reduced = static_cast<float*>(
pool.get_symmetric(Names::kDxReduced, g_reduced_bytes));
workspace.sync = static_cast<std::uint64_t*>(
pool.get_symmetric(Names::kDxSync, sync_bytes));
workspace.launch_epoch = static_cast<const std::uint64_t*>(
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
liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/workspace.cuh
+147
View File
@@ -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 <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>
#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<float> 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
liger_cute_kernels/csrc/core/src/nvshmem.cu
+3
View File
@@ -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;
});
liger_cute_kernels/liger_cute_kernels/tvm_ffi.py
+120
View File
@@ -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
liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp
+263
View File
@@ -2,11 +2,18 @@
#include <tvm/ffi/tvm_ffi.h>
#include <cuda_runtime.h>
#include <nccl.h>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <exception>
#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<int64_t*>(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<int64_t>(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<int64_t>(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<int>(nranks), id, static_cast<int>(rank)),
"ncclCommInitRank");
static_cast<int64_t*>(out.data_ptr())[0] =
reinterpret_cast<int64_t>(comm);
}
void nccl_comm_destroy(int64_t comm_handle) {
if (comm_handle == 0) {
return;
}
CheckNccl(
ncclCommDestroy(reinterpret_cast<ncclComm_t>(comm_handle)),
"ncclCommDestroy");
}
void my_pe(ffi::TensorView out) {
RequireCpuInt32(out, 1);
CheckStatus(liger_cute_nvshmem_my_pe(static_cast<int*>(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<int>(max_tokens), static_cast<int>(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<int>(max_tokens), static_cast<int>(max_hidden),
static_cast<int>(max_local_vocab), static_cast<int>(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<const int64_t*>(target.data_ptr());
params.gemm.tokens = static_cast<int>(tokens);
params.gemm.hidden = static_cast<int>(hidden);
params.gemm.local_vocab = static_cast<int>(local_vocab);
params.gemm.vocab_start = vocab_start;
params.gemm.ignore_index = ignore_index;
params.gemm.inverse_temperature = static_cast<float>(inverse_temperature);
params.nll = static_cast<float*>(nll.data_ptr());
params.lse = static_cast<float*>(lse.data_ptr());
params.entropy = static_cast<float*>(entropy.data_ptr());
params.nccl_comm_handle = nccl_comm_handle;
cudaStream_t stream = reinterpret_cast<cudaStream_t>(
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<true, 90>(params, stream);
} else {
liger::fused_scaled_linear_cross_entropy::
fused_linear_scaled_cross_entropy_forward<false, 90>(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<cudaStream_t>(
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<int>(local_vocab));
BackwardTpParamsSm90<90> params;
params.gemm.x = x.data_ptr();
params.gemm.weight = weight.data_ptr();
params.gemm.target = static_cast<const int64_t*>(target.data_ptr());
params.gemm.grad_output = static_cast<const float*>(grad_output.data_ptr());
params.gemm.lse = static_cast<const float*>(lse.data_ptr());
params.gemm.entropy = static_cast<const float*>(entropy.data_ptr());
params.gemm.entropy_grad = static_cast<const float*>(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<int>(tokens);
params.gemm.hidden = static_cast<int>(hidden);
params.gemm.local_vocab = static_cast<int>(local_vocab);
params.gemm.vocab_start = vocab_start;
params.gemm.ignore_index = ignore_index;
params.gemm.inverse_temperature = static_cast<float>(inverse_temperature);
params.tiles_per_reduce = static_cast<int>(tiles_per_reduce);
params.num_comm_channels = 1;
if (return_entropy) {
liger::fused_scaled_linear_cross_entropy::
fused_linear_scaled_cross_entropy_backward<true, 90>(params, stream);
} else {
liger::fused_scaled_linear_cross_entropy::
fused_linear_scaled_cross_entropy_backward<false, 90>(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);
liger_cute_kernels/test/test_fslce_teams.py
+282
View File
@@ -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)
liger_cute_kernels/tests/cpp/CMakeLists.txt
+14
View File
@@ -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
liger_cute_kernels/tests/cpp/test_fused_scaled_linear_cross_entropy.cu
+195
View File
@@ -0,0 +1,195 @@
#include <gtest/gtest.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>
#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<BackwardConfig, 4, 2>;
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<float>) == 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<CommConfig>(5), 3);
fslce::DxTileGroup first =
fslce::make_dx_tile_group<CommConfig>(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<CommConfig>(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<true>(1.0f, true, false);
rank0.fold<true>(3.0f, true, false);
fslce::OnlineSoftmaxState rank1;
rank1.fold<true>(2.0f, true, true);
rank1.fold<true>(4.0f, true, false);
float global_max = fmaxf(rank0.max_value, rank1.max_value);
fslce::CorrectedSoftmaxStats corrected0 =
fslce::correct_for_global_max<true>(rank0, global_max);
fslce::CorrectedSoftmaxStats corrected1 =
fslce::correct_for_global_max<true>(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<true>(
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<true>(-2.0f, true, false);
lhs.fold<true>(0.5f, true, false);
fslce::OnlineSoftmaxState rhs;
rhs.fold<true>(3.0f, true, true);
rhs.fold<true>(1.0f, true, false);
fslce::OnlineSoftmaxState merged =
fslce::merge_online_softmax<true>(lhs, rhs);
fslce::OnlineSoftmaxState direct;
direct.fold<true>(-2.0f, true, false);
direct.fold<true>(0.5f, true, false);
direct.fold<true>(3.0f, true, true);
direct.fold<true>(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<false>(2.0f, true, false);
state.fold<false>(4.0f, true, false);
EXPECT_FLOAT_EQ(state.exp_weighted_sum, 0.0f);
}
TEST(TensorParallelFusedScaledLinearCrossEntropyOnlineSoftmax, ZerosIgnoredEntropy) {
fslce::FinalizedSoftmax result = fslce::finalize_softmax<true>(
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();
}