[FlyDSL] Add optional ROCm/AMD FlyDSL backend with cross-entropy and fused linear cross-entropy

[kashif]

Summary

Adds an optional FlyDSL backend for AMD/ROCm, starting with cross-entropy and fused linear cross-entropy. FlyDSL is a Python DSL + MLIR compiler targeting CDNA/RDNA, so this gives ROCm users a non-Triton path for the two ops that matter most for LLM training memory.

It's strictly opt-in — nothing changes unless you ask for it:

pip install -e ".[flydsl]"
LIGER_KERNEL_IMPL=flydsl python train.py

Ops that aren't implemented fall back to the default Triton kernels, and unsupported options (class weights, token accuracy, predicted tokens, token scaling) raise NotImplementedError rather than silently diverging.

Design notes

Two things worth calling out, because they're where most of the performance lives:

Normalizers stay on the device. n_non_ignore is passed to the kernel as a small fp32 tensor rather than a host float, so the mean reduction is folded into the loss and the gradient inside the kernel. No .item(), no D2H sync, and no rescale pass over the full (BT, V) tensor afterwards. As a side effect the stored per-row losses stay at ~L/N instead of ~L, which is what keeps sum(loss_1d) from overflowing fp16 at large batch sizes.

The gradient is produced in backward, not forward. Forward saves only the per-row lse (BT floats); backward rebuilds softmax(x) = exp2((x - lse) * log2e) and folds grad_output straight into the coefficient. That means we never materialize an unscaled gradient, so there's no second pass to rescale it. Total traffic is 2 reads + 1 write of the logits — the same as the Triton kernel's online softmax — with no host sync anywhere.

The kernel itself follows the FlyDSL norm/softmax kernels: online softmax (one sweep for max+sum), 128-bit buffer copies over the whole row with a single masked tail vector, and a shuffle_xor butterfly block reduction. Indexing by vector element rather than by tile matters in practice — real vocab sizes (32000, 128256, 152064) aren't multiples of BLOCK * vec_width, and gating on that would drop every production shape onto a scalar, one-element-per-lane path.

FLCE chunks over the batch dimension using the same aspect_ratio rule PyTorch's LinearCrossEntropyOptions uses (next_pow2(ceil(BT / ceil(V/H)))), so peak logits memory is chunk×V rather than BT×V.

Benchmarks

AMD Radeon 890M (gfx1150, RDNA 3.5, wave32), ROCm 7.2, bf16. Timed with FlyDSL's own do_bench (CUDA events, warmup + median); peak memory measured in isolated subprocesses.

Cross-entropy, fwd+bwd — vs the Triton kernel:

BT × V	FlyDSL	Triton	ratio
4096 × 32768	9.08 ms	5.97 ms	1.52×
4096 × 32000	8.87 ms	5.80 ms	1.53×
2048 × 128256	22.10 ms	17.68 ms	1.25×
2048 × 152064	29.17 ms	20.92 ms	1.39×
2048 × 2048	0.54 ms	0.49 ms	1.10×

Fused linear CE — BT=8192, H=2048, V=128256 (Llama-3 lm_head), full logits would be 2101 MB:

variant	peak memory	time
naive F.cross_entropy(lm_head(h))	12642 MB	2410 ms
torch.nn.functional.linear_cross_entropy	1153 MB	4803 ms
Liger FLCE (Triton)	2236 MB	2849 ms
Liger FLCE (FlyDSL)	1712 MB	2203 ms

So on this hardware FLCE lands 24% below Triton on memory and 23% faster, and beats PyTorch's chunked linear_cross_entropy on speed by 2.2× while using more memory than it.

CE is still ~1.2–1.5× behind Triton on raw kernel time. It's latency-bound rather than bandwidth-bound (cutting HBM traffic 33% via register-buffering the row changed nothing), so the next lever is occupancy/scheduling — I'd want an rocprofv3 ATT trace before guessing further. Happy to land that as a follow-up.

Tests

132 tests across test/transformers/test_flydsl_cross_entropy.py and test_flydsl_fused_linear_cross_entropy.py, all passing on gfx1150. They check parity against the Triton kernel across dtypes (fp32/fp16/bf16), reductions (mean/sum/none), ignore_index, label_smoothing, softcap, z_loss, non-contiguous inputs, and the all-ignored edge case; plus parity against F.linear_cross_entropy, the chunk-size heuristic, and a memory invariant asserting FLCE never allocates the full BT×V logits.

Everything skips cleanly when FlyDSL isn't installed or there's no AMD GPU, so this is a no-op for existing CI.

[kashif]

While reviewing this I found that the reduction="none" gradient bug I inherited here is not specific to FlyDSL, it is in the Triton backend on main too. Filed separately as #1298.

Short version: with reduction="none" the loss is per-token, so backward gets a per-token grad_output. But grad_weight is a sum over tokens, so the upstream weight has to be folded into each row of the logits gradient before that sum. element_mul_kernel loads grad_output as a scalar and scales the already-summed result, which cannot be correct. The existing tests only ever pass torch.ones, where the scaling is a no-op, so it went unnoticed.

I have pushed the equivalent fix to this branch (facf21f) so the FlyDSL backend does not ship with the same bug: for reduction="none" the gradient is deferred to backward and grad_output is folded in per row before projecting. mean/sum keep the existing compute-in-forward fast path. Added a regression test with a non-uniform grad_output, which fails without the fix and passes with it.

#1298 is independent of this PR and can land on its own.