[apullin]

Summary

Fixes four related problems in the fused DyT kernels:

• Launch limit: both kernels placed rows on CUDA grid dimension 1, which holds at most 65,535 blocks, so more than 65,535 rows failed to launch. Both kernels now use a 1D grid.
• Small widths: one program per row starves the GPU when rows have only 16-128 channels. Widths up to 256 columns now process several rows per 4,096-element tile (exposed as a BLOCK_M autotune dimension, alongside the existing BLOCK_N autotune), and the backward kernel runs up to four persistent programs per SM on that path. Wider rows keep one row per program.
• Padding: masked loads alone do not keep padded rows out of the reductions. With infinite alpha, a padded zero gives tanh(inf * 0) = NaN and poisons the parameter gradients, so padded contributions are now masked after the arithmetic in the backward kernel.
• Precision: dalpha and dbeta were cast to x.dtype, unlike dgamma. They now take their parameters' dtypes, so float32 parameter gradients are not rounded through e.g. bfloat16. The output keeps the x.dtype contract.

Because the autotuned BLOCK_N is not known before launch, the backward da/dg/db buffers are sized for the worst-case config (most programs, smallest BLOCK_N) and zeroed, so programs and column blocks a config does not use contribute nothing to the reductions. The reset_to_zero=["DA"] autotune behaviour is preserved.

Testing Done

• Hardware Type: none available to the author (no GPU on this host)
• run make checkstyle equivalent: ruff check and ruff format --check pass on both changed files
• run make test to ensure correctness
• run make test-convergence to ensure convergence

The new and extended tests in test/transformers/test_dyt.py (small/ragged widths, 65,537 rows, ragged rows with infinite alpha, and bfloat16 input with float32 parameters) still need to run on a GPU, e.g. via CI. This is opened as a draft until those results are in.

[apullin]

Closing this draft for now: the branch remains pushed on my fork, and the PR message has not been decided yet. It can be reopened, or a new PR opened from the same branch, once the message is settled.

[apullin]

Temporary GPU-testing note, kept in sync with GPU_TESTING.md on the branch (commit c991f18). That file must be removed before this is opened as a real PR — it is not part of the change.

GPU TESTING — TEMPORARY, DO NOT MERGE

Remove this file before opening the pull request. It is not part of
the change. Delete it with:

git rm GPU_TESTING.md
git commit --amend --no-edit   # or a follow-up commit, before pushing

Then verify it is gone with git status and git show --stat HEAD.

This branch has not been executed on a GPU yet. The port was reworked
for the upstream autotune kernels (BLOCK_M as an autotune dimension,
backward buffers sized for the worst-case autotuned config), so it must
be tested on a real GPU before the pull request leaves draft. The
results below are what the pull request's Testing Done section will
report.

Proposed pull request message (draft, not yet posted)

Title: fix(dyt): use a 1D grid and tile small-width rows

## Summary

Modification to the shape handling in the fused DyT kernels that gives
a major performance gain in my application: a small network with DyT,
as tested on an H100.

- Both kernels placed rows on CUDA grid dimension 1, which holds at most
  65,535 blocks, so more than 65,535 rows failed to launch. Both kernels
  now use a 1D grid.
- One program per row starves the GPU on narrow rows. Widths up to 256
  columns now process several rows per 4,096-element tile, exposed as a
  `BLOCK_M` autotune dimension; wider rows keep one row per program, and
  the backward kernel runs up to four persistent programs per SM on the
  row-tiled path.
- Padded rows are now masked out after the arithmetic in the backward
  kernel: with infinite `alpha`, `tanh(inf * 0)` is NaN and masked loads
  alone let it poison the parameter gradients.
- `dalpha` and `dbeta` now take their parameters' dtypes, like `dgamma`
  already did, instead of rounding float32 parameter gradients through
  `x.dtype`. The output keeps the `x.dtype` contract.

Because the autotuned `BLOCK_N` is not known before launch, the backward
buffers are sized for the worst-case config and zeroed, so unused
programs/column blocks contribute nothing to the reductions.

## Testing Done

Hardware: H100.

- `make checkstyle`: passes.
- `make test`: [to be filled in from the rented-box run]
- `make test-convergence`: [to be filled in from the rented-box run]

Runbook: rented H100 box

Run this in a disposable copy. make checkstyle can auto-fix/format,
so check git status afterwards (apart from this file, it should be
clean).

git clone --branch dyt-1d-grid-row-tiling \
  https://github.com/apullin/Liger-Kernel.git liger-pr
cd liger-pr
git rev-parse HEAD   # record this hash with the results
nvidia-smi
python3 --version    # 3.10-3.12 preferred

python3 -m venv .venv && source .venv/bin/activate
pip install --upgrade pip
pip install -e ".[dev]"

make checkstyle 2>&1 | tee /tmp/checkstyle.log

# Gate: the DyT tests. The first run per hidden size triggers autotune
# benchmarking across all configs — let it run, do not kill it as hung.
# This includes the 65,537-row case.
python -m pytest test/transformers/test_dyt.py -v 2>&1 | tee /tmp/test_dyt.log

If rental time allows, then the full template set:

make test 2>&1 | tee /tmp/make_test.log
make test-convergence 2>&1 | tee /tmp/test_convergence.log

Note on convergence: no convergence mini-model uses DyT, so it should
not be affected by this change. Treat the DyT tests plus make test
as the real gate, and convergence as template completeness if the box
is still up.

Report back

• nvidia-smi output, and torch / triton versions.
• The commit hash tested (git rev-parse HEAD).
• Pass/fail counts and the tails of /tmp/checkstyle.log,
  /tmp/test_dyt.log, and (if run) /tmp/make_test.log and
  /tmp/test_convergence.log.
• Any failures, verbatim.
• git status (should be clean apart from this file, if present).

Do not mark the pull request ready for review until these results are
in. Remove this file first.