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

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,14 +167,15 @@ 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,
         alpha,
         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
 
 

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