Add SM100 tensor-parallel fused scaled cross entropy by kolehma8 · Pull Request #1458 · linkedin/Liger-Kernel

kolehma8 · 2026-09-10T20:33:17Z

Summary

Adds a native Blackwell-family (sm_100f, including B200/B300) tensor-parallel fused scaled linear cross-entropy implementation.

Fuses the classifier GEMM with NLL/LSE and optional entropy.
Adds a persistent SM100 backward that executes dZ, dX, dW, token-wave scheduling, node-local NVLS, and two-host reduction in one cluster-2 launch.
Supports TP1/2/4/8 node-local execution and two-host hierarchical TP with 1/2/4/8 selected GPUs per host, including TP subgroups of a larger NVSHMEM world.
Uses the shared forward-verified QP put-signal ring for inter-host communication.
Distributes the remote FP32 dX merge over warp 1 from every resident CTA, synchronized with an HBM atomic completion counter.
Adds reusable workspace configuration/query APIs, public API documentation, correctness validators, phase benchmarks, and topology contract tests.

Code entry points

Layer	Forward	Backward
PyTorch autograd	`LigerFusedLinearScaledCrossEntropyNativeTPFunction.forward`	`LigerFusedLinearScaledCrossEntropyNativeTPFunction.backward`
TVM FFI	`fused_linear_scaled_cross_entropy_forward`	`fused_linear_scaled_cross_entropy_backward`
SM100 host launcher	`fused_linear_scaled_cross_entropy_forward_sm100`	`fused_linear_scaled_cross_entropy_backward_sm100`
Device kernel	`forward_gemm_tp_kernel_sm100`	`backward_gemm_tp_kernel_sm100`
Role implementations	`forward_gemm_roles_sm100.cuh`	`backward_gemm_mainloop_sm100.cuh`

Multi-host subgroup topology

Multi-host setup is relative to the configured TP team rather than the full NVSHMEM world. nvshmem_team_split_2d derives:

one local NVLS row containing only the selected GPUs on the current host;
one matching-rank remote column containing the corresponding selected GPU from each host.

For a 16-rank world on two eight-GPU hosts, TP group {0,4,8,12} becomes local rows {0,4} / {8,12} and remote pairs {0,8} / {4,12}. SM100 remote kernels use a parent-team stream barrier followed by a cooperative clustered CUDA launch, so disjoint TP groups do not require world-synchronous kernel entry.

Validated concurrently as four TP4 groups {0,4,8,12}, {1,5,9,13}, {2,6,10,14}, and {3,7,11,15}:

28/28 numerical cases pass across the four groups.
Covers entropy, ignored targets, ragged vocabulary, full waves, and two-wave tails.
Cross-rank dX equality remains exact.
Deliberately staggering group entry by 500 ms also passes for forward and backward, confirming subgroup-scoped progress.
Two-host local widths 1, 2, 4, and 8 GPUs per host pass.
NVSHMEM_IB_ENABLE_IBGDA=1 was forced for the final campaign. Runtime initialization reported Successfully initialized the transport: IBGDA. It will be used for device-side APIs over IB, created DCI/RC QPs, and selected each configured mlx5_* device.

Random BF16, M=4096, H=4096, global V=131072, 10 warmups, and 30 measured iterations:

TP group root	Native forward ms	Native forward TF/GPU	Forward speedup	Native F+B ms	Native F+B TF/GPU	F+B speedup
0	1.080	1017.9	2.82x	5.193	846.9	1.69x
1	1.072	1025.8	2.85x	5.166	851.3	1.69x
2	1.058	1039.5	2.86x	5.128	857.6	1.69x
3	1.090	1009.2	2.78x	5.175	849.9	1.68x

Native versus fallback forward maximum absolute difference was 5.45e-4.

SM90 regression validation

Validated the shared topology, workspace, and binding changes on two eight-GPU
H200 hosts with a clean sm_90a build:

Complete native core build passed.
12/12 shared reduction and parent-relative topology tests passed.
8/8 SM90 FSLCE configuration and online-softmax contract tests passed.
12/12 standalone SM90 forward GEMM numerical cases passed.
Public FFI world-team and process-group subteam tests passed, including
forward, dX, dW, entropy, and CUDA graph replay.
A two-host TP16 world team passed forward, dX, dW, and entropy validation.
Four concurrent strided TP4 groups {0,4,8,12}, {1,5,9,13},
{2,6,10,14}, and {3,7,11,15} all passed on H200.
NVSHMEM reported IBGDA as the active device-side InfiniBand transport on all
16 ranks.

Implementation

Forward

2SM UMMA, M256 x N256 x K64, cluster 2 x 1.
Five-stage TMA mainloop.
Four source slots and a default wave width of 64 N256 vocabulary tiles.
Warp 0 performs local reduction, warp 1 runs the matching-rank host ring, warp 2 owns TMA, warp 3 issues UMMA, and warps 4-11 drain TMEM/epilogue work.
The __global__ entry point is isolated in forward_gemm_kernel_sm100.cuh and dispatches those roles explicitly; templated implementations live in forward_gemm_roles_sm100.cuh.

Backward

One persistent 384-thread cluster-2 kernel.
dZ: five-stage M256 x N256 x K64.
dX: four-stage 2SM N512 schedule, publishing two M128xN256 communication tiles.
dW: six-stage paired-N schedule with TMA store for wave 0 and TMA reduce-add for later waves.
A single 512-column Allocator2Sm TMEM allocation is reused across all phases.
TP16:
1. every CTA's warp 0 performs node-local NVLS reduce-scatter;
2. CTA 0 warp 1 performs the existing inter-host ring transfer;
3. warp 1 from all resident CTAs merges a global-strided portion of the remote FP32 shard;
4. one HBM atomic arrival per CTA releases CTA 0;
5. warp 0 performs node-local all-gather and BF16 scatter.

The measured TP16 remote stage is 0.285-0.328 ms:

Component	Latency
IBRC transfer, 8 MiB/GPU	0.119-0.156 ms
All-CTA FP32 merge	0.160-0.176 ms
Total remote stage	0.285-0.328 ms
In-kernel dW overlap window	~2.41 ms
Exposed remote tail	0 ms

Correctness

TP1/2/4/8/16: 35/35 distributed numerical cases pass.
Covers aligned and ragged shapes, entropy, ignored targets, 4K waves, 4097-token boundaries, and 5000-token two-wave tails.
Cross-rank dX equality is exact.
Maximum observed relative error:
- dX: 0.262%
- dW: 0.498%
Forward native versus Verl maximum absolute difference: 1.09e-3.
C++ tests:
- shared reduction/topology: 11/11
- SM100 backward contracts: 12/12
make checkstyle passes.

Performance methodology

GPU: NVIDIA B300.
BF16 inputs, FP32 accumulation.
Hidden size: 4096.
Fixed global vocabulary: 131072.
Tokens: 4096, 8192, 16384.
TP: 1, 2, 4, 8, 16.
10 warmups, 30 measured iterations, median CUDA-event latency.
Verl baseline: repository 512-token chunked TP fallback.
Forward+backward throughput counts four classifier GEMMs: forward logits,
backward logits recomputation for dZ, dX, and dW.
Efficiency is per-GPU TFLOP/s relative to TP1 for the same shape.
TP16 uses two eight-GPU hosts and explicit one-HCA-per-local-GPU mapping.

M=4096 forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	2.852	1542.0	100.0%	8.489	518.1	100.0%	2.98x
2	1.457	1508.9	97.9%	4.653	472.6	91.2%	3.19x
4	0.764	1438.3	93.3%	2.856	384.9	74.3%	3.74x
8	0.428	1283.6	83.2%	2.134	257.6	49.7%	4.98x
16	0.465	591.0	38.3%	2.096	131.1	25.3%	4.51x

M=4096 forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	10.950	1606.5	100.0%	29.898	588.4	100.0%	2.73x
2	6.170	1425.6	88.7%	15.570	564.9	96.0%	2.52x
4	3.562	1234.7	76.9%	8.451	520.4	88.5%	2.37x
8	2.253	976.0	60.8%	5.426	405.2	68.9%	2.41x
16	2.171	506.5	31.5%	4.472	245.9	41.8%	2.06x

M=8192 forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	5.744	1531.4	100.0%	17.107	514.2	100.0%	2.98x
2	2.867	1533.8	100.2%	9.376	469.1	91.2%	3.27x
4	1.469	1496.9	97.7%	5.638	390.1	75.9%	3.84x
8	0.793	1385.8	90.5%	4.185	262.7	51.1%	5.28x
16	0.762	721.6	47.1%	4.077	134.8	26.2%	5.35x

M=8192 forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	22.264	1580.3	100.0%	58.375	602.8	100.0%	2.62x
2	11.600	1516.5	96.0%	31.618	556.4	92.3%	2.73x
4	6.463	1360.9	86.1%	17.348	507.1	84.1%	2.68x
8	4.219	1042.5	66.0%	10.653	412.8	68.5%	2.52x
16	4.020	546.9	34.6%	8.659	254.0	42.1%	2.15x

M=16384 forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	12.645	1391.2	100.0%	35.020	502.4	100.0%	2.77x
2	6.399	1374.5	98.8%	19.560	449.7	89.5%	3.06x
4	3.206	1371.9	98.6%	11.426	384.9	76.6%	3.56x
8	1.674	1313.3	94.4%	8.197	268.3	53.4%	4.90x
16	1.449	759.0	54.6%	9.809	112.1	22.3%	6.77x

M=16384 forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	50.387	1396.5	100.0%	115.214	610.8	100.0%	2.29x
2	23.270	1512.0	108.3%	62.505	562.9	92.2%	2.69x
4	12.579	1398.5	100.1%	35.498	495.6	81.1%	2.82x
8	8.281	1062.1	76.1%	21.188	415.2	68.0%	2.56x
16	7.898	556.8	39.9%	16.706	263.2	43.1%	2.12x

Memory

Inputs, weights, and final outputs are excluded. Native reports its exact reusable C++/NVSHMEM forward+backward workspace. Verl reports incremental PyTorch peak memory.

M=4096

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	1810.40	1024.10	2720.11
2	1202.42	512.10	1376.11
4	898.46	256.10	704.11
8	746.53	128.10	368.11
16	682.67	64.10	200.11

M=8192

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	2004.79	1024.19	2752.21
2	1300.81	512.19	1408.21
4	948.84	256.19	736.21
8	772.91	128.19	400.21
16	709.06	64.19	232.21

M=16384

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	2393.56	1024.35	2816.40
2	1497.58	512.35	1472.40
4	1049.61	256.35	800.40
8	825.68	128.35	464.40
16	761.83	64.35	296.40

Additional weak-scaling production result

With fixed V_local=65536, M=4096, and H=4096:

TP	Native forward TF/GPU	Verl forward TF/GPU	Native F+B TF/GPU	Verl F+B TF/GPU
1	1524.5	488.6	1607.1	570.3
2	1505.3	471.8	1424.4	560.7
4	1513.3	459.2	1541.9	543.2
8	1490.8	446.8	1580.4	537.9
16	1323.0	412.2	1434.5	533.7

Model-shape campaign

The following production model shapes use M=8192, random BF16 inputs, 10 warmups, and 30 measured iterations. TP16 spans the same two eight-GPU hosts as the fixed-shape campaign.

Qwen3-30B-A3B (`H=2048`, `V=151936`)

Forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	3.692	1380.8	100.0%	15.885	320.9	100.0%	4.30x
2	1.937	1315.9	95.3%	9.192	277.3	86.4%	4.75x
4	1.040	1225.4	88.7%	5.581	228.4	71.2%	5.37x
8	0.592	1076.2	77.9%	4.136	154.1	48.0%	6.99x
16	0.664	479.8	34.7%	4.012	79.4	24.7%	6.04x

Forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	13.692	1489.3	100.0%	47.689	427.6	100.0%	3.48x
2	7.480	1363.2	91.5%	26.363	386.8	90.4%	3.52x
4	4.100	1243.3	83.5%	14.834	343.7	80.4%	3.62x
8	2.507	1016.5	68.3%	9.720	262.3	61.3%	3.88x
16	2.545	500.8	33.6%	8.346	152.7	35.7%	3.28x

Memory

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	1975.79	1187.19	1962.08
2	1286.31	593.69	996.65
4	941.84	296.94	515.33
8	769.41	152.19	276.94
16	695.56	74.37	154.58

Llama 4 (`H=5120`, `V=202048`)

Forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	11.727	1445.3	100.0%	28.506	594.6	100.0%	2.43x
2	5.795	1462.3	101.2%	15.228	556.5	93.6%	2.63x
4	2.880	1471.2	101.8%	8.489	499.1	83.9%	2.95x
8	1.479	1433.0	99.1%	5.339	396.8	66.7%	3.61x
16	1.103	960.3	66.4%	7.578	139.8	23.5%	6.87x

Forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	50.180	1351.1	100.0%	100.330	675.7	100.0%	2.00x
2	22.677	1494.8	110.6%	54.966	616.7	91.3%	2.42x
4	11.785	1438.3	106.4%	29.663	571.5	84.6%	2.52x
8	6.939	1221.3	90.4%	16.230	522.1	77.3%	2.34x
16	6.026	703.2	52.0%	23.633	179.3	26.5%	3.92x

Memory

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	2655.29	1578.69	5015.46
2	1626.31	792.19	2547.99
4	1111.84	394.81	1315.52
8	854.41	200.19	698.15
16	756.06	98.84	389.87

DeepSeek-V3 (`H=7168`, `V=129280`)

Forward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	10.923	1390.0	100.0%	22.352	679.2	100.0%	2.05x
2	5.317	1427.6	102.7%	11.923	636.7	93.7%	2.24x
4	2.638	1439.0	103.5%	6.943	546.7	80.5%	2.63x
8	1.356	1399.8	100.7%	4.556	416.5	61.3%	3.36x
16	1.041	911.7	65.6%	4.862	195.2	28.7%	4.67x

Forward + backward

TP	Native ms	Native TF/GPU	Efficiency	Verl ms	Verl TF/GPU	Efficiency	Speedup
1	43.588	1393.3	100.0%	81.751	742.9	100.0%	1.88x
2	22.167	1369.9	98.3%	44.494	682.4	91.9%	2.01x
4	12.411	1223.3	87.8%	24.582	617.6	83.1%	1.98x
8	7.600	998.8	71.7%	13.475	563.3	75.8%	1.77x
16	7.024	540.4	38.8%	10.843	350.0	47.1%	1.54x

Memory

TP	Native workspace MiB	Verl forward peak MiB	Verl F+B peak MiB
1	2278.79	1010.19	4279.46
2	1437.81	505.19	2196.71
4	1017.34	256.19	1156.21
8	807.41	128.19	634.08
16	744.56	64.19	373.04

Validation commands

make checkstyle

On B300 (-DLIGER_CUTE_CUDA_ARCH=100f):

cmake --build build --target \
  libliger_cute_kernels.so \
  test_fslce_shared_reduction \
  test_fused_scaled_linear_cross_entropy_backward_sm100

./build/tests/cpp/test_fslce_shared_reduction
./build/tests/cpp/test_fused_scaled_linear_cross_entropy_backward_sm100

Distributed numerical validation:

python -m torch.distributed.run --standalone --nproc_per_node=<1|2|4|8> \
  liger_cute_kernels/tests/python/validate_fused_scaled_linear_cross_entropy_backward_sm100.py

TP16 uses two eight-GPU hosts with:

export NVSHMEM_REMOTE_TRANSPORT=ibrc
export NVSHMEM_IB_ENABLE_IBGDA=1
export NVSHMEM_ENABLE_NIC_PE_MAPPING=1
unset NVSHMEM_HCA_LIST
export NVSHMEM_HCA_PE_MAPPING='mlx5_0:1:1,mlx5_2:1:1,mlx5_3:1:1,mlx5_4:1:1,mlx5_5:1:1,mlx5_6:1:1,mlx5_8:1:1,mlx5_9:1:1'

yueyiming2009

Public API and workspace review comments. Both findings are blocking for the advertised SM100 integration; kernel review is still in progress.

Implement persistent SM100 forward and backward kernels with NVLS and two-host remote-ring reduction, add topology-aware workspace management and public APIs, and include correctness and performance coverage through TP16. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Split the forward kernel entry point from its role implementations, dispatch warp roles explicitly at the kernel top level, and consolidate repeated dW accumulator handling while addressing review feedback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Derive parent-relative local and matching-rank teams with a 2D NVSHMEM split, use subgroup-scoped cooperative SM100 launches, and validate strided groups such as {0,4,8,12} across two hosts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Record the production IBGDA environment and transport initialization signal for multi-host TP subgroup execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Count the backward logits recomputation alongside forward, dX, and dW when reporting end-to-end throughput. Refresh the README headline with the corrected four-GEMM rates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Normalize reference qualifiers in the compile-time output contract so NVCC accepts the SM90 test while preserving the intended exact type check. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Dispatch the native TP implementation on Hopper and SM100-family GPUs using the real CUDA capability query. Reject workspace-size requests that do not exactly match an existing immutable configuration, preventing mixed configured/requested footprints. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

vaibhavjindal

LGTM! Just want to check if the verl convergence tests pass and also if the test suite present here is equivalent to regular FLCE test suite that we have.

kolehma8 · 2026-09-14T22:19:11Z

LGTM! Just want to check if the verl convergence tests pass and also if the test suite present here is equivalent to regular FLCE test suite that we have.

I will follow-up this PR with a Verl training run once the integration is finished.

kolehma8 commented Sep 10, 2026

View reviewed changes

yueyiming2009 reviewed Sep 12, 2026

View reviewed changes

Comment thread src/liger_kernel/ops/fused_linear_scaled_cross_entropy.py

Comment thread liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp

jkolehma and others added 7 commits September 14, 2026 10:00

Document IBGDA subgroup validation

2925532

Record the production IBGDA environment and transport initialization signal for multi-host TP subgroup execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Fix fused cross entropy FLOP accounting

542cc51

Count the backward logits recomputation alongside forward, dX, and dW when reporting end-to-end throughput. Refresh the README headline with the corrected four-GEMM rates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Fix SM90 forward output type assertion

56fae70

Normalize reference qualifiers in the compile-time output contract so NVCC accepts the SM90 test while preserving the intended exact type check. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kolehma8 force-pushed the jkolehm/tp-fslce-sm100 branch from c727e07 to f6683fe Compare September 14, 2026 17:30

vaibhavjindal approved these changes Sep 14, 2026

View reviewed changes

kolehma8 added this pull request to the merge queue Sep 14, 2026

Merged via the queue into linkedin:main with commit f2ac965 Sep 14, 2026
1 of 3 checks passed

kolehma8 deleted the jkolehm/tp-fslce-sm100 branch September 14, 2026 22:19

Navigation Menu

Add SM100 tensor-parallel fused scaled cross entropy - #1458

Add SM100 tensor-parallel fused scaled cross entropy#1458
kolehma8 merged 7 commits into
linkedin:mainfrom
kolehma8:jkolehm/tp-fslce-sm100

kolehma8 commented Sep 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yueyiming2009 left a comment

Uh oh!

Uh oh!

Uh oh!

vaibhavjindal left a comment

Uh oh!

kolehma8 commented Sep 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

kolehma8 commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Code entry points

Multi-host subgroup topology

SM90 regression validation

Implementation

Forward

Backward

Correctness

Performance methodology

M=4096 forward

M=4096 forward + backward

M=8192 forward

M=8192 forward + backward

M=16384 forward

M=16384 forward + backward

Memory

M=4096

M=8192

M=16384

Additional weak-scaling production result

Model-shape campaign

Qwen3-30B-A3B (H=2048, V=151936)

Forward

Forward + backward

Memory

Llama 4 (H=5120, V=202048)

Forward

Forward + backward

Memory

DeepSeek-V3 (H=7168, V=129280)

Forward

Forward + backward

Memory

Validation commands

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yueyiming2009 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

vaibhavjindal left a comment

Choose a reason for hiding this comment

Uh oh!

kolehma8 commented Sep 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kolehma8 commented Sep 10, 2026 •

edited

Loading

Qwen3-30B-A3B (`H=2048`, `V=151936`)

Llama 4 (`H=5120`, `V=202048`)

DeepSeek-V3 (`H=7168`, `V=129280`)