[bhargavchintam]

Fixes #1605

Summary

liger_dyt_fwd launches _dyt_fwd_kernel with grid = (cdiv(N, BLOCK_N), M), which puts the row count M (batch × sequence length) on grid dimension 1. CUDA limits that dimension to 65535 (the B200 I ran this on reports that limit), so any input with more than 65535 rows fails at launch with Triton Error [CUDA]: invalid argument. A batch of 128 sequences of 1024 tokens is 131072 rows.

This PR swaps the two grid dimensions in the forward launch and the two tl.program_id calls in _dyt_fwd_kernel, so the rows are on dimension 0 and the column tiles on dimension 1. The tiled SwiGLU kernel already orders its grid this way (grid = (n_rows, triton.cdiv(n_cols, BLOCK_SIZE)) in ops/swiglu.py). The backward kernel is unchanged because its dimension 1 is the SM count.

Details

• CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X/Y/Z on the B200 are 2147483647, 65535 and 65535.
• ops/backends/_ascend/ops/dyt.py is not changed, and I did not run anything on Ascend.
• The swap changes the order in which programs are launched. Forward time on a B200 (my own script: liger_dyt_fwd launched 300 times after a warmup, with torch.cuda.synchronize() around the loop, mean in µs):

shape	dtype	main	this PR
(8, 512, 256)	fp32	24.2	23.8
(8, 512, 256)	bf16	23.6	23.9
(16, 1024, 2048)	fp32	41.1	41.7
(16, 1024, 2048)	bf16	38.3	37.8
(32, 256, 4096)	fp32	41.1	43.6
(32, 256, 4096)	bf16	38.3	37.9
(8, 128, 16384)	fp32	25.5	26.7
(8, 128, 16384)	bf16	25.3	24.9

Testing Done

• Hardware Type: NVIDIA B200 (NGC PyTorch 26.10 container: torch 2.15.0a0, Triton 3.8.0, CUDA 13.4)
• run make test to ensure correctness
• run make checkstyle to ensure code style (I ran its two check-only steps, shown at the end, because the other two steps rewrite files)
• run make test-convergence to ensure convergence

I did not run make test or make test-convergence. The B200 was rented for this run, and on it I ran only the changed kernel's test file, test/transformers/test_dyt.py, as below.

New test test_liger_dyt_more_rows_than_grid_dimension_limit uses 2 × 32769 = 65538 rows, with and without beta, and compares the output and all gradients with torch.

$ python3 -m pytest test/transformers/test_dyt.py -k more_rows_than_grid --disable-warnings -q --tb=line

# main (e44de705) with only the new test applied
E   RuntimeError: Triton Error [CUDA]: invalid argument
E   RuntimeError: Triton Error [CUDA]: invalid argument
FAILED test/transformers/test_dyt.py::test_liger_dyt_more_rows_than_grid_dimension_limit[False]
FAILED test/transformers/test_dyt.py::test_liger_dyt_more_rows_than_grid_dimension_limit[True]
================= 2 failed, 24 deselected, 2 warnings in 1.55s =================

# this PR
test/transformers/test_dyt.py::test_liger_dyt_more_rows_than_grid_dimension_limit[False] PASSED [ 50%]
test/transformers/test_dyt.py::test_liger_dyt_more_rows_than_grid_dimension_limit[True] PASSED [100%]
================= 2 passed, 24 deselected, 1 warning in 16.61s =================

$ python3 -m pytest test/transformers/test_dyt.py --disable-warnings -q --tb=line
================== 26 passed, 1 warning in 140.93s (0:02:20) ===================

The snippet from the issue, run on this branch (on main the second line is RuntimeError: Triton Error [CUDA]: invalid argument):

torch.Size([65535, 8])
torch.Size([65536, 8])

$ ruff check --output-format=concise .
All checks passed!

$ ruff format --check --diff .
500 files already formatted

[bhargavchintam]

I ran make test on a B200 with this PR and #1608 applied together on main (e44de705; torch 2.15.0a0, Triton 3.8.0, transformers 5.22.2):

= 22 failed, 5702 passed, 1054 skipped, 21 xfailed, 318 warnings, 6 rerun in 2820.14s (0:47:00) =

None of the 22 is in test_dyt.py or test_group_norm.py. I re-ran those 22 on plain e44de705 and on this tree. The same 21 fail on both, and the other one, a DPO loss test, passes on both when run that way. I did not look into why those 21 fail, and I did not run make test-convergence.