From 4acae82cf4c2fe18899f386f93589f56f017df49 Mon Sep 17 00:00:00 2001 From: Bindu Bhargava R Chintam Date: Sun, 8 Nov 2026 23:30:17 -0800 Subject: [PATCH] fix(dyt): put the row count on the first grid dimension of the forward kernel CUDA limits grid dimensions 1 and 2 to 65535 programs. The forward kernel put the row count (batch * sequence length) on dimension 1, so any input with more than 65535 rows failed at launch with "Triton Error [CUDA]: invalid argument". A batch of 128 sequences of 1024 tokens is 131072 rows. Swap the two dimensions, as the tiled SwiGLU kernel already does. The column tile count stays on dimension 1, and the backward kernel keeps its grid because it uses the SM count there. --- src/liger_kernel/ops/dyt.py | 7 ++++--- test/transformers/test_dyt.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/src/liger_kernel/ops/dyt.py b/src/liger_kernel/ops/dyt.py index ed5c74886..05b6252dd 100644 --- a/src/liger_kernel/ops/dyt.py +++ b/src/liger_kernel/ops/dyt.py @@ -32,9 +32,9 @@ ) @triton.jit def _dyt_fwd_kernel(X, Y, Alpha, Gamma, Beta, HAVE_BETA: tl.constexpr, N: tl.constexpr, BLOCK_N: tl.constexpr): - col = tl.cast(tl.program_id(0), tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) + col = tl.cast(tl.program_id(1), tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) mask = col < N - row_id = tl.cast(tl.program_id(1), tl.int64) + row_id = tl.cast(tl.program_id(0), tl.int64) X += row_id * N Y += row_id * N @@ -107,7 +107,8 @@ def liger_dyt_fwd(x, alpha, gamma, beta): y = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M) + # CUDA limits grid dimensions 1 and 2 to 65535 programs, so the row count has to go on dimension 0. + grid = lambda meta: (M, triton.cdiv(N, meta["BLOCK_N"])) _dyt_fwd_kernel[grid]( x, y, diff --git a/test/transformers/test_dyt.py b/test/transformers/test_dyt.py index 2f28e672d..882dac480 100644 --- a/test/transformers/test_dyt.py +++ b/test/transformers/test_dyt.py @@ -100,6 +100,36 @@ def test_liger_dyt_correctness(B, T, hidden_size, beta, init_alpha, dtype, atol, ) +@pytest.mark.parametrize("beta", [False, True]) +def test_liger_dyt_more_rows_than_grid_dimension_limit(beta): + # CUDA limits grid dimensions 1 and 2 to 65535 programs, and this input has 65538 rows. + B, T, hidden_size = 2, 32769, 8 + _input = torch.randn(B, T, hidden_size, device=device, dtype=torch.float32) + x1 = _input.clone().requires_grad_(True) + x2 = _input.clone().requires_grad_(True) + + torch_dyt = TorchDyT(hidden_size=hidden_size, beta=beta).to(device) + liger_dyt = LigerDyT(hidden_size=hidden_size, beta=beta).to(device) + + torch_output = torch_dyt(x1) + liger_output = liger_dyt(x2) + assert_verbose_allclose(torch_output, liger_output, rtol=1e-5, atol=1e-5, extra_info="[output]") + + grad_output = torch.randn_like(_input) + torch_output.backward(grad_output) + liger_output.backward(grad_output) + assert_verbose_allclose(x1.grad, x2.grad, rtol=1e-5, atol=1e-5, extra_info="[input.grad]") + + # These gradients are sums over all the rows, so they carry more fp32 rounding error. + rtol, atol = 1e-3, 1e-3 + assert_verbose_allclose(torch_dyt.alpha.grad, liger_dyt.alpha.grad, rtol=rtol, atol=atol, extra_info="[alpha.grad]") + assert_verbose_allclose(torch_dyt.gamma.grad, liger_dyt.gamma.grad, rtol=rtol, atol=atol, extra_info="[gamma.grad]") + if beta: + assert_verbose_allclose( + torch_dyt.beta.grad, liger_dyt.beta.grad, rtol=rtol, atol=atol, extra_info="[beta.grad]" + ) + + @pytest.mark.parametrize("beta", [False, True]) @pytest.mark.parametrize( "B, T, hidden_size",