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))