From 2524c6edb0d07e7ee6a31ac94c6712a40aaa5b96 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Wed, 11 Nov 2026 14:45:30 -0800 Subject: [PATCH] fix(dyt): use a 1D grid and tile small-width rows The forward and backward 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. Widths up to 256 columns process several rows per 4,096-element tile as an autotune dimension, instead of one program per row, which starves the GPU on narrow rows. 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. --- src/liger_kernel/ops/dyt.py | 186 +++++++++++++++++++++++----------- test/transformers/test_dyt.py | 42 ++++++++ 2 files changed, 170 insertions(+), 58 deletions(-) diff --git a/src/liger_kernel/ops/dyt.py b/src/liger_kernel/ops/dyt.py index ed5c74886..fb9f1e73b 100644 --- a/src/liger_kernel/ops/dyt.py +++ b/src/liger_kernel/ops/dyt.py @@ -10,6 +10,40 @@ from liger_kernel.ops.utils import infer_device from liger_kernel.utils import is_npu_available +# Feature widths whose block is at most this many columns get BLOCK_M rows per +# program (millions of rows of 16-128 channels, where one program per row +# starves the GPU). Wider rows keep one row per program. +_ROW_TILE_MAX_BLOCK_N = 256 +# Elements per program on the row-tiled path, the block size used by +# `_block_rms_norm` for the same tradeoff. +_ROW_TILE_ELEMENTS = 4096 +# Backward programs per SM on the row-tiled path. +_ROW_TILE_BWD_PROGRAMS_PER_SM = 4 +# Smallest BLOCK_N in the autotune configs below. The backward buffers are +# sized for this worst case because the autotuned BLOCK_N is not known +# before launch (see liger_dyt_bwd). +_MIN_AUTOTUNE_BLOCK_N = 64 + + +def _dyt_configs(): + configs = [] + for block_n in (64, 128, 256, 512, 1024, 2048, 4096): + block_m = _ROW_TILE_ELEMENTS // block_n if block_n <= _ROW_TILE_MAX_BLOCK_N else 1 + warp_counts = (4, 8) if block_n < 1024 else (4, 8, 16) + for num_stages in (1, 2): + for num_warps in warp_counts: + configs.append( + triton.Config({"BLOCK_N": block_n, "BLOCK_M": block_m}, num_stages=num_stages, num_warps=num_warps) + ) + return configs + + +def _dyt_bwd_num_programs(meta, num_rows, num_sms): + if meta["BLOCK_M"] > 1: + return min(num_sms * _ROW_TILE_BWD_PROGRAMS_PER_SM, triton.cdiv(num_rows, meta["BLOCK_M"])) + return num_sms + + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: # typical import path with dispatch available @@ -22,80 +56,106 @@ @triton.autotune( - configs=[ - triton.Config({"BLOCK_N": bn}, num_stages=ns, num_warps=nw) - for bn in [1024, 2048, 4096] - for ns in [1, 2] - for nw in [4, 8, 16] - ], + configs=_dyt_configs(), key=["N"], ) @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) - mask = col < N - row_id = tl.cast(tl.program_id(1), tl.int64) +def _dyt_fwd_kernel( + X, + Y, + Alpha, + Gamma, + Beta, + HAVE_BETA: tl.constexpr, + M, + N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # A 1D grid, which allows 2**31 - 1 programs (dimensions 1 and 2 allow + # 65,535), with a row's column blocks adjacent. + pid = tl.cast(tl.program_id(0), tl.int64) + num_col_blocks = tl.cdiv(N, BLOCK_N) + row = (pid // num_col_blocks) * BLOCK_M + tl.arange(0, BLOCK_M) + col = (pid % num_col_blocks) * BLOCK_N + tl.arange(0, BLOCK_N) + col_mask = col < N + mask = (row < M)[:, None] & col_mask[None, :] + offsets = row[:, None] * N + col[None, :] - X += row_id * N - Y += row_id * N alpha = tl.load(Alpha).to(tl.float32) - - gamma = tl.load(Gamma + col, mask=mask, other=0.0).to(tl.float32) - - x = tl.load(X + col, mask=mask, other=0.0).to(tl.float32) + gamma = tl.load(Gamma + col, mask=col_mask, other=0.0).to(tl.float32) + x = tl.load(X + offsets, mask=mask, other=0.0).to(tl.float32) tanh_x = tanh(alpha * x) - y = tanh_x * gamma + y = tanh_x * gamma[None, :] if HAVE_BETA: - beta = tl.load(Beta + col, mask=mask, other=0.0).to(tl.float32) - y += beta - tl.store(Y + col, y, mask=mask) + beta = tl.load(Beta + col, mask=col_mask, other=0.0).to(tl.float32) + y += beta[None, :] + tl.store(Y + offsets, y, mask=mask) @triton.autotune( - configs=[ - triton.Config({"BLOCK_N": bn}, num_stages=ns, num_warps=nw) - for bn in [1024, 2048, 4096] - for ns in [1, 2] - for nw in [4, 8, 16] - ], + configs=_dyt_configs(), key=["N"], - # DA is indexed by program_id(0), so different BLOCK_N configs write to - # different slot counts per SM. Autotune trials don't zero outputs between - # runs, so stale slots from a prior trial would leak into da.sum(). Reset - # DA between trials to isolate each config's writes. - reset_to_zero=["DA"], + # DA, DG, and DB are sized for the worst-case config, so different + # (BLOCK_M, BLOCK_N) configs write to different slot/row counts. Autotune + # trials don't zero outputs between runs, so stale values from a prior + # trial would leak into the sums in liger_dyt_bwd. Reset them between + # trials to isolate each config's writes. + reset_to_zero=["DA", "DG", "DB"], ) @triton.jit def _dyt_bwd_kernel( - DY, DX, DA, DG, DB, X, Alpha, Gamma, HAVE_BETA: tl.constexpr, M, N: tl.constexpr, BLOCK_N: tl.constexpr + DY, + DX, + DA, + DG, + DB, + X, + Alpha, + Gamma, + HAVE_BETA: tl.constexpr, + M, + N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, ): - col = tl.cast(tl.program_id(0), tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) - mask = col < N - start_row_id = tl.cast(tl.program_id(1), tl.int64) + pid = tl.cast(tl.program_id(0), tl.int64) + num_col_blocks = tl.cdiv(N, BLOCK_N) + row_block_id = pid // num_col_blocks + col_block_id = pid % num_col_blocks + col = col_block_id * BLOCK_N + tl.arange(0, BLOCK_N) + col_mask = col < N alpha = tl.load(Alpha).to(tl.float32) - da = 0.0 - gamma = tl.load(Gamma + col, mask=mask, other=0.0).to(tl.float32) - dg = tl.zeros((BLOCK_N,), dtype=tl.float32) + gamma = tl.load(Gamma + col, mask=col_mask, other=0.0).to(tl.float32) + da = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + dg = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) if HAVE_BETA: - db = tl.zeros((BLOCK_N,), dtype=tl.float32) - for row_id in range(start_row_id, M, tl.num_programs(1)): - x = tl.load(X + row_id * N + col, mask=mask, other=0.0).to(tl.float32) - dy = tl.load(DY + row_id * N + col, mask=mask, other=0.0).to(tl.float32) + db = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + num_row_programs = tl.num_programs(0) // num_col_blocks + for start_row in range(row_block_id * BLOCK_M, M, num_row_programs * BLOCK_M): + row = start_row + tl.arange(0, BLOCK_M) + # Mask padded rows out after the arithmetic, not just on load: + # with infinite alpha, tanh(inf * 0) is NaN and would poison the + # parameter gradients if the padded rows were reduced. + mask = (row < M)[:, None] & col_mask[None, :] + offsets = row[:, None] * N + col[None, :] + x = tl.load(X + offsets, mask=mask, other=0.0).to(tl.float32) + dy = tl.load(DY + offsets, mask=mask, other=0.0).to(tl.float32) tanh_x = tanh(alpha * x) if HAVE_BETA: - db += dy - dg += dy * tanh_x - tmp = (1 - tanh_x * tanh_x) * dy * gamma - da += tl.sum(x * tmp, 0) + db += tl.where(mask, dy, 0.0) + dg += tl.where(mask, dy * tanh_x, 0.0) + tmp = (1 - tanh_x * tanh_x) * dy * gamma[None, :] + da += tl.where(mask, x * tmp, 0.0) dx = alpha * tmp - tl.store(DX + row_id * N + col, dx, mask=mask) + tl.store(DX + offsets, dx, mask=mask) - tl.store(DG + start_row_id * N + col, dg, mask=mask) + tl.store(DG + row_block_id * N + col, tl.sum(dg, 0), mask=col_mask) if HAVE_BETA: - tl.store(DB + start_row_id * N + col, db, mask=mask) - tl.store(DA + start_row_id * tl.cdiv(N, 512) + tl.program_id(0), da) + tl.store(DB + row_block_id * N + col, tl.sum(db, 0), mask=col_mask) + tl.store(DA + pid, tl.sum(tl.sum(da, 1), 0)) def liger_dyt_fwd(x, alpha, gamma, beta): @@ -107,7 +167,7 @@ def liger_dyt_fwd(x, alpha, gamma, beta): y = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M) + grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),) _dyt_fwd_kernel[grid]( x, y, @@ -115,6 +175,7 @@ def liger_dyt_fwd(x, alpha, gamma, beta): gamma, beta, HAVE_BETA, + M, N, ) return y.view(input_shape) @@ -134,17 +195,26 @@ def liger_dyt_bwd(dy, x, alpha, gamma, beta): NUM_SMS = torch.xpu.get_device_properties(x.device).gpu_subslice_count elif device == "npu": NUM_SMS = get_npu_core_count() - da = torch.zeros(NUM_SMS, triton.cdiv(N, 512), dtype=torch.float32, device=x.device) - dg = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) - db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None + # Sized for the worst-case autotuned config (most programs, smallest + # BLOCK_N), and zeroed so programs/column blocks a config does not use + # contribute nothing to the sums below. + max_programs = NUM_SMS * _ROW_TILE_BWD_PROGRAMS_PER_SM + da = torch.zeros(max_programs, triton.cdiv(N, _MIN_AUTOTUNE_BLOCK_N), dtype=torch.float32, device=x.device) + dg = torch.zeros(max_programs, N, dtype=torch.float32, device=x.device) + db = ( + torch.zeros(max_programs, N, dtype=torch.float32, device=x.device) + if HAVE_BETA + else torch.empty(0, dtype=torch.float32, device=x.device) + ) dx = torch.empty_like(dy) - grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS) + grid = lambda meta: (_dyt_bwd_num_programs(meta, M, NUM_SMS) * triton.cdiv(N, meta["BLOCK_N"]),) _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N) - if HAVE_BETA: - db = db.sum(0).to(x.dtype) + db = db.sum(0).to(beta.dtype) if HAVE_BETA else None dg = dg.sum(0).to(gamma.dtype) - da = da.sum().to(x.dtype).unsqueeze(0) + # dalpha takes alpha's dtype, like dgamma takes gamma's: rounding a + # float32 parameter gradient through x.dtype (e.g. bfloat16) loses it. + da = da.sum().to(alpha.dtype).unsqueeze(0) return dx.view(input_shape), da, dg, db diff --git a/test/transformers/test_dyt.py b/test/transformers/test_dyt.py index 2f28e672d..427b421b6 100644 --- a/test/transformers/test_dyt.py +++ b/test/transformers/test_dyt.py @@ -49,6 +49,9 @@ def forward(self, x): (4, 16, 2048), (1, 1, 1023), # Minimal batch/seq with near power-of-2 hidden (3, 7, 256), # Prime numbers for batch/seq + (2, 4099, 16), # Small hidden, several rows per program + (1, 1001, 61), # Small odd hidden, ragged row blocks + (1, 65537, 16), # More rows than CUDA grid dimension 1 allows ], ) @pytest.mark.parametrize( @@ -108,6 +111,9 @@ def test_liger_dyt_correctness(B, T, hidden_size, beta, init_alpha, dtype, atol, (4, 16, 2048), (1, 1, 1023), # Minimal batch/seq with near power-of-2 hidden (3, 7, 256), # Prime numbers for batch/seq + (2, 4099, 16), # Small hidden, several rows per program + (1, 1001, 61), # Small odd hidden, ragged row blocks + (1, 65537, 16), # More rows than CUDA grid dimension 1 allows ], ) @pytest.mark.parametrize( @@ -158,3 +164,39 @@ def test_liger_dyt_functional(B, T, hidden_size, beta, dtype, atol, rtol): assert_verbose_allclose(gamma1.grad, gamma2.grad, rtol=rtol, atol=atol) if beta: assert_verbose_allclose(beta1.grad, beta2.grad, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("rows", [1, 257]) +def test_liger_dyt_ragged_rows_with_infinite_alpha(rows): + # 257 rows tile raggedly when several rows share a program. Padded rows + # must not poison the parameter gradients: tanh(inf * 0) is NaN. + x = torch.ones(rows, 16, device=device) + alpha = torch.full((1,), float("inf"), device=device, requires_grad=True) + gamma = torch.ones(16, device=device, requires_grad=True) + beta = torch.zeros(16, device=device, requires_grad=True) + output = LigerDyTFunction.apply(x, alpha, gamma, beta) + dalpha, dgamma, dbeta = torch.autograd.grad(output, (alpha, gamma, beta), torch.ones_like(output)) + assert_verbose_allclose(dalpha, torch.zeros_like(alpha)) + assert_verbose_allclose(dgamma, torch.full_like(gamma, rows)) + assert_verbose_allclose(dbeta, torch.full_like(beta, rows)) + + +@pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU") +def test_liger_dyt_bfloat16_input_float32_parameters(): + # 257 is not representable in bfloat16 (rounds to 256), so rounding + # dalpha or dbeta through x.dtype loses integer precision on float32 + # parameters. + rows, channels = 257, 16 + x = torch.ones(rows, channels, device=device, dtype=torch.bfloat16, requires_grad=True) + alpha = torch.zeros(1, device=device, dtype=torch.float32, requires_grad=True) + gamma = torch.ones(channels, device=device, dtype=torch.float32, requires_grad=True) + beta = torch.zeros(channels, device=device, dtype=torch.float32, requires_grad=True) + output = LigerDyTFunction.apply(x, alpha, gamma, beta) + assert output.dtype == torch.bfloat16 + dalpha, dgamma, dbeta = torch.autograd.grad(output, (alpha, gamma, beta), torch.ones_like(output)) + assert dalpha.dtype == torch.float32 + assert dgamma.dtype == torch.float32 + assert dbeta.dtype == torch.float32 + assert_verbose_allclose(dalpha, torch.full_like(alpha, rows * channels)) + assert_verbose_allclose(dgamma, torch.zeros_like(gamma)) + assert_verbose_allclose(dbeta, torch.full_like(beta, rows))