[CUDA] Enable native CuTe DSL fused linear cross entropy on SM103

[heiheiha798]

Summary

• Enable the existing native CuTe DSL fused linear cross entropy path on exact SM103 in addition to SM100.
• Resolve architecture and capability from the actual input tensor device; exact (blackwell, 10.0) and (blackwell_ultra, 10.3) pairs are accepted, while SM120, Hopper, CPU, FP32, and reduction="none" retain their existing behavior.
• Hold the actual input device context across the complete native forward, including the CE stage and retained-graph recomputation, so non-current-device launches acquire the correct stream and architecture metadata.
• Extend native-hardware, non-default-device, compile-cache, DLPack ABI, non-contiguous floating-input, and full-pipeline device-context coverage. The persistent GEMM implementation and its SM100 schedule are unchanged.

This change tracks fork issue #17. The original OPEN draft engineering record is heiheiha798/Liger-Kernel#18.

This is complementary to #1396: that PR adds Triton fallback behavior for capability misses, while this PR enables the existing native path on exact SM103.

Testing Done

• Hardware Type: NVIDIA B300 SXM6 AC (SM103)
• run make test to ensure correctness (the focused suites listed below were run instead)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence (the selected mini-Llama3 and mini-Qwen2 cases listed below were run instead)

All GPU validation ran through Slurm on one NVIDIA B300 SXM6 AC (SM103), driver 580.126.09, CUDA 13.0, PyTorch 2.13.0+cu130, Triton 3.7.1, NVIDIA CUTLASS DSL 4.7.0.

• Job 5085, profile-first gate: unchanged native run_epilogue_gemm passed BF16/FP16 PyTorch GEMM oracles; unchanged native public FLCE passed independent FP32 PyTorch loss and all-gradient oracles for BF16/FP16, mean/sum, bias, ignore rows, label smoothing, and non-unit upstream gradients.
• Job 5089, clean exact production commit 990b24c: 247 CuTe FLCE/GEMM tests, 203 transformer/monkey-patch tests, and BF16 mini-Llama3/mini-Qwen2 selection (4 selected cases) passed.
• Job 5093: real DLPack ABI and BF16/FP16 non-contiguous floating-input cases passed.
• Job 5095, clean commit c95cdaf: 250 CuTe FLCE/GEMM tests passed, including TVM-FFI, DLPack, partial/frozen gradients, repeated backward, ignored rows, odd tails, actual-device dispatch, and compile-cache isolation.
• Job 5102, clean exact PR head a1db3ac: the focused full-pipeline device-context regression passed and the complete CuTe FLCE/GEMM suite passed (251 passed). The test forces the native pipeline with mocked launch bodies and proves its CE stage executes inside the input-device context even when the modeled ambient device differs.
• make checkstyle passed on the login node at final head: 370 files lint-clean and already formatted.

B300 Results

Job 5089 used detached clean worktrees at baseline 124fb8a and production commit 990b24c, fresh exact-source processes, identical deterministic inputs, compilation excluded, three measurements per round, and seven alternating paired rounds. Public forward-plus-backward is authoritative. The token sweep fixes H4096/V128256; the vocabulary sweep fixes BT8192/H4096.

Sweep	Shape	Triton ms	CuTe ms	Speedup	Incremental peak allocation ratio
tokens	BT 1024	115.786	4.012	28.862x	3.544x
tokens	BT 2048	116.371	6.279	18.534x	3.534x
tokens	BT 4096	118.188	9.907	11.929x	3.170x
tokens	BT 8192	120.851	20.232	5.975x	3.141x
tokens	BT 16384	132.159	41.617	3.176x	3.087x
tokens	BT 32768	163.901	80.394	2.042x	2.993x
tokens	BT 65536	227.261	163.709	1.390x	2.846x
vocab	V 32000	10.572	4.802	2.202x	2.992x
vocab	V 102400	53.379	15.629	3.414x	3.162x
vocab	V 128256	120.877	20.040	6.033x	3.141x
vocab	V 152064	143.475	25.191	5.697x	3.157x
vocab	V 201088	189.034	33.465	5.649x	3.178x
vocab	V 262144	477.293	43.516	10.969x	3.171x

All 7/7 token and 6/6 vocabulary points exceeded the 20% acceptance target. Thirteen FP16, odd-H/V, feature, metric, token-scaling, and FP32-gradient-accumulation controls had no median regression above 5%; the minimum ratio was 0.969x on the sub-millisecond V513 control. At BT8192/H4096/V128256, incremental peak allocation fell from 4,134.7 MiB to 1,316.5 MiB, a 2,818.2 MiB reduction. Job 5102 reconfirmed the final context-only head at the same reference point: 120.922 ms Triton versus 20.268 ms native CuTe (5.966x), with 3.140x lower incremental peak allocation.

NCU from job 5089 captured one unchanged native persistent forward GEMM chunk at BT8192/H4096/V128256: 862.592 us, grid (2,1,74), block (192,1,1), 48 registers/thread, 94.95% tensor-pipe activity, 1.089 GB read, 249.553 MB written, 1.552 TB/s aggregate DRAM traffic, 72.60% L2 hit rate, and zero local loads/stores or spill requests. The public native path launches 8 persistent GEMMs for the reference forward chunking.

Reproduction

Run the existing repository benchmark from clean detached baseline/candidate worktrees in the same B300 allocation, using the same environment and model selection:

cd /path/to/baseline
LIGER_KERNEL_IMPL= python benchmark/scripts/benchmark_fused_linear_cross_entropy.py --model llama_3_8b
cd /path/to/candidate
LIGER_KERNEL_IMPL=cutedsl python benchmark/scripts/benchmark_fused_linear_cross_entropy.py --model llama_3_8b

Raw local evidence: .cache/b300-cutedsl-sm103/final-5089.out, final-5089.err, native-5089.ncu-rep, native-5089.csv, head-5095.out, review-5102.out, and review-5102.err. Job 5088 failed before worker import because its metadata probe omitted the isolated dependency hook. Jobs 5090/5092 stopped in newly added test setup before FLCE; corrected job 5093 passed. These discarded runs are not performance evidence.

Optional SM103 schedule tuning was intentionally omitted: the unchanged native schedule cleared every frozen public-path gate, avoiding architecture-specific schedule churn.