[wqh-creator]

Summary

Fix FP16/BF16 Sparsemax input gradients by subtracting the support mean in FP32 before the final output cast. The previous backward kernel rounded that mean to the output dtype first, losing small differences between upstream gradients.

For a zero input row with upstream BF16 gradients [1, 1.0078125], the correct input gradient is [-0.00390625, 0.00390625]. Current main returns [0, 0.0078125], incorrectly producing a nonzero row sum. FP16 has the same failure with adjacent representable upstream values. The forward path is unchanged.

Add exact regression coverage for FP16/BF16, last/non-last dimensions, and an inactive support entry. The tests compare against the existing PyTorch reference computed in FP32 and assert the zero-sum gradient invariant.

Duplicate check: searched open and merged Sparsemax PRs/issues. #1545 fixes forward shift invariance; #1497 fixes row offsets; #1459 changes cuTile contracts. None fixes this premature cast in the native Triton backward.

AI assistance: this change, tests and GPU experiments were prepared with OpenAI Codex.

Testing Done

• Hardware Type: NVIDIA GeForce RTX 5090 (SM120, 32 GB), CUDA driver 570.144; Python 3.12.12, torch 2.10.0+cu128, Triton 3.6.0.
• Base: e44de705ecc9b43eb33aa35ad10ee9d59a054c3b.
• Before the fix: CUDA_VISIBLE_DEVICES=4 python -m pytest test/transformers/test_sparsemax.py -k preserves_small_gradient_differences -v → 4 failed, 10 deselected, with the expected gradient mismatches.
• After the fix: CUDA_VISIBLE_DEVICES=4 python -m pytest test/transformers/test_sparsemax.py -xvs → 14 passed, including the large-row-offset test; no skips.
• run make test to ensure correctness — the scoped Sparsemax suite above was run; the full repository suite was not run.
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence — not run; no model patching or forward behavior is changed.

Bounded benchmarks reuse benchmark/scripts/benchmark_sparsemax.py::setup_sparsemax and the repository speed/memory helpers for BF16 shapes 128×1024 and 1024×4096. Forward, backward and full modes were measured with three repeats per configuration, including alternating base/patched processes. Full-mode peak allocations were unchanged: 2.508 MiB and 80.031 MiB respectively.

Small-shape eager full-mode timings were variable: approximately 0.045 ms on the base and 0.045–0.175 ms on the patch, despite isolated backward timings remaining around 0.006 ms. This host-inclusive variability has not been conclusively explained, so this PR makes no performance improvement claim. Additional fixed-input CUDA Graph full forward/backward measurements (two alternating runs per version, five measurements each) did not show a corresponding GPU execution regression:

BF16 shape	Base median / range (ms)	Patched median / range (ms)
128×1024	0.02372 / 0.02211–0.02585	0.02361 / 0.02291–0.02448
1024×4096	0.14677 / 0.14598–0.14860	0.14574 / 0.14291–0.14910

Correctness is the purpose of this change. Model-training convergence and throughput are not claimed.