diff --git a/fla/ops/kda/chunk_inter.py b/fla/ops/kda/chunk_inter.py index 994d2c164b..d85b2a1cf0 100644 --- a/fla/ops/kda/chunk_inter.py +++ b/fla/ops/kda/chunk_inter.py @@ -6,8 +6,8 @@ import triton.language as tl from fla.ops.utils import prepare_chunk_indices -from fla.ops.utils.op import exp -from fla.utils import autotune_cache_kwargs, check_shared_mem +from fla.ops.utils.op import exp, make_tensor_descriptor +from fla.utils import autotune_cache_kwargs, check_shared_mem, is_tma_supported BK_LIST = [32, 64] if check_shared_mem() else [16, 32] BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] @@ -51,6 +51,7 @@ def chunk_kda_bwd_kernel_inter( BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + USE_TMA: tl.constexpr, IS_VARLEN: tl.constexpr, ): i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -91,29 +92,46 @@ def chunk_kda_bwd_kernel_inter( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + if not USE_TMA: + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + else: + desc_v = make_tensor_descriptor(v, [T, V], [H*V, 1], [BT, BV]) + desc_do = make_tensor_descriptor(do, [T, V], [H*V, 1], [BT, BV]) + desc_h = make_tensor_descriptor(h, [V, K], [1, V], [BV, BK]) + desc_dh = make_tensor_descriptor(dh, [V, K], [1, V], [BV, BK]) + desc_dv = make_tensor_descriptor(dv, [T, V], [H*V, 1], [BT, BV]) + # [BT, BV] + b_v = desc_v.load([i_t * BT, i_v * BV]) + b_do = desc_do.load([i_t * BT, i_v * BV]) + b_dv = desc_dv.load([i_t * BT, i_v * BV]) + # [BV, BK] + b_h = desc_h.load([i_v * BV, i_k * BK]) + b_dh = desc_dh.load([i_v * BV, i_k * BK]) # [BK] b_dgk += tl.sum(b_h * b_dh, axis=0) # [BT, BK] b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) - p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + if not USE_TMA: + p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + else: + desc_dw = make_tensor_descriptor(dw, [T, K], [H*K, 1], [BT, BK]) + desc_dw.store([i_t * BT, i_k * BK], -b_dw.to(b_dw.dtype)) b_dgk *= exp(b_gn) b_dq *= scale @@ -184,5 +202,6 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) K=K, V=V, BT=BT, + USE_TMA=is_tma_supported, ) return dq, dk, dw, dg diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py index 806b0cfe5b..011c63bbf5 100644 --- a/fla/ops/kda/chunk_intra.py +++ b/fla/ops/kda/chunk_intra.py @@ -223,6 +223,7 @@ def chunk_kda_bwd_kernel_intra( BC: tl.constexpr, BK: tl.constexpr, NC: tl.constexpr, + USE_TMA: tl.constexpr, IS_VARLEN: tl.constexpr, ): i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -269,17 +270,30 @@ def chunk_kda_bwd_kernel_intra( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + if not USE_TMA: + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + else: + desc_k = make_tensor_descriptor(k, [T, K], [H*K, 1], [BC, BK]) + desc_g = make_tensor_descriptor(g, [T, K], [H*K, 1], [BC, BK]) + desc_dAqk = make_tensor_descriptor(dAqk, [T, BT], [H*BT, 1], [BC, BC]) + desc_dAkk = make_tensor_descriptor(dAkk, [T, BT], [H*BT, 1], [BC, BC]) + # [BC, BK] + b_k = desc_k.load([i_t * BT + i_j * BC, i_k * BK]) + b_gk = desc_g.load([i_t * BT + i_j * BC, i_k * BK]) + # [BC, BC] + b_dAqk = desc_dAqk.load([i_t * BT + i_i * BC, i_j * BC]) + b_dAkk = desc_dAkk.load([i_t * BT + i_i * BC, i_j * BC]) + b_kg = b_k * exp(b_gn[None, :] - b_gk) - # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) # [BC, BK] b_dq2 += tl.dot(b_dAqk, b_kg) b_dk2 += tl.dot(b_dAkk, b_kg) @@ -292,10 +306,16 @@ def chunk_kda_bwd_kernel_intra( p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + if not USE_TMA: + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + else: + desc_q = make_tensor_descriptor(q, [T, K], [H*K, 1], [BC, BK]) + desc_k = make_tensor_descriptor(k, [T, K], [H*K, 1], [BC, BK]) + b_q = desc_q.load([i_t * BT + i_i * BC, i_k * BK]) + b_k = desc_k.load([i_t * BT + i_i * BC, i_k * BK]) for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC] @@ -315,16 +335,25 @@ def chunk_kda_bwd_kernel_intra( b_db = tl.sum(b_dk2 * b_k, 1) b_dk2 *= b_b[:, None] - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + if not USE_TMA: + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_dq = tl.load(p_dq, boundary_check=(0, 1)) + else: + desc_dq = make_tensor_descriptor(dq, [T, K], [H*K, 1], [BC, BK]) + desc_dq2 = make_tensor_descriptor(dq2, [T, K], [H*K, 1], [BC, BK]) + b_dq = desc_dq.load([i_t * BT + i_i * BC, i_k * BK]) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) b_dg = b_q * b_dq2 - b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) - tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) + b_dq2 = b_dq2 + b_dq tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + if not USE_TMA: + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) + else: + desc_dq2.store([i_t * BT + i_i * BC, i_k * BK], b_dq2.to(b_dq.dtype)) - tl.debug_barrier() b_dkt = tl.zeros([BC, BK], dtype=tl.float32) NC = min(NC, tl.cdiv(T - i_t * BT, BC)) @@ -333,27 +362,42 @@ def chunk_kda_bwd_kernel_intra( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) # [BC] b_b = tl.load(p_b, boundary_check=(0,)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + if not USE_TMA: + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + else: + desc_q = make_tensor_descriptor(q, [T, K], [H*K, 1], [BC, BK]) + desc_k = make_tensor_descriptor(k, [T, K], [H*K, 1], [BC, BK]) + desc_g = make_tensor_descriptor(g, [T, K], [H*K, 1], [BC, BK]) + desc_dAqk = make_tensor_descriptor(dAqk, [BT, T], [1, H*BT], [BC, BC]) + desc_dAkk = make_tensor_descriptor(dAkk, [BT, T], [1, H*BT], [BC, BC]) + # [BC, BK] + b_q = desc_q.load([i_t * BT + i_j * BC, i_k * BK]) + b_kb = desc_k.load([i_t * BT + i_j * BC, i_k * BK]) * b_b[:, None] + b_gk = desc_g.load([i_t * BT + i_j * BC, i_k * BK]) + # [BC, BC] + b_dAqk = desc_dAqk.load([i_i * BC, i_t * BT + i_j * BC]) + b_dAkk = desc_dAkk.load([i_i * BC, i_t * BT + i_j * BC]) o_j = i_t * BT + i_j * BC + o_i m_j = o_j < T # [BC, BK] - b_qg = b_q * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) - b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) + b_gkn = tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) + b_qg = b_q * b_gkn + b_kbg = b_kb * b_gkn # [BC, BK] # (SY 09/17) important to not use bf16 here to have a good precision. b_dkt += tl.dot(b_dAqk, b_qg) @@ -510,11 +554,12 @@ def chunk_kda_bwd_intra( B, T, H, K = k.shape BT = chunk_size BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) + BK = min(32, triton.next_power_of_2(K)) if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + # NC = 4 NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) @@ -546,6 +591,7 @@ def chunk_kda_bwd_intra( BC=BC, BK=BK, NC=NC, + USE_TMA=is_tma_supported, ) dq = dq2 dk = dk2 diff --git a/tests/ops/test_kda.py b/tests/ops/test_kda.py index 0e2bb1eb85..aa10fb55fe 100644 --- a/tests/ops/test_kda.py +++ b/tests/ops/test_kda.py @@ -136,21 +136,21 @@ def test_fused_recurrent( @pytest.mark.parametrize( - ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype', 'tma'), + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype', 'tma', 'triltf32'), [ pytest.param( *test, - id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}-tma{}".format(*test), + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}-tma{}-triltf32{}".format(*test), ) for test in [ - (1, 63, 1, 64, 1, 1, 0, False, torch.float16, True), - (2, 500, 3, 60, 1, 1, 0, False, torch.float16, True), - (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16, False), - (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16, False), - (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16, True), - (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16, True), - (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16, False), - (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16, True), + (1, 63, 1, 64, 1, 1, 0, False, torch.float16, True, True), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16, True, False), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16, False, True), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16, False, False), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16, True, True), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16, True, True), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16, False, True), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16, True, True), ] ], ) @@ -165,12 +165,17 @@ def test_chunk( use_qk_l2norm_in_kernel: bool, dtype: torch.dtype, tma: bool, + triltf32: bool, ): torch.manual_seed(42) if not tma: os.environ['FLA_USE_TMA'] = '0' else: os.environ['FLA_USE_TMA'] = '1' + if triltf32: + os.environ['FLA_TRIL_PRECISION'] = 'tf32x3' + else: + os.environ['FLA_TRIL_PRECISION'] = 'ieee' q = torch.rand(B, T, H, D, dtype=dtype) k = torch.rand(B, T, H, D, dtype=dtype) v = torch.rand(B, T, H, D, dtype=dtype) @@ -219,18 +224,19 @@ def test_chunk( assert_close('dg', ref_dg, tri_dg, 0.02) assert_close('db', ref_db, tri_db, 0.02) assert_close('dh0', ref_dh0, tri_dh0, 0.008) - + os.environ['FLA_USE_TMA'] = '0' + os.environ['FLA_TRIL_PRECISION'] = 'ieee' @pytest.mark.parametrize( - ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype', 'use_tma'), + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype', 'tma', 'triltf32'), [ - pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-tma{}".format(*test)) + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-tma{}-triltf32{}".format(*test)) for test in [ - (4, 60, 0, [0, 15], torch.float16, True), - (4, 64, 0, [0, 256, 500, 1000], torch.float16, True), - (4, 128, 0.5, [0, 256, 500, 1000], torch.float16, False), - (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16, True), - (4, 256, 0, [0, 15, 100, 300, 1200, 4096], torch.float16, False), + (4, 60, 0, [0, 15], torch.float16, True, False), + (4, 64, 0, [0, 256, 500, 1000], torch.float16, True, False), + (4, 128, 0.5, [0, 256, 500, 1000], torch.float16, False, True), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16, True, True), + (4, 256, 0, [0, 15, 100, 300, 1200, 4096], torch.float16, False, True), ] ], ) @@ -244,12 +250,17 @@ def test_chunk_varlen( mask_p: float, cu_seqlens: list[int], dtype: torch.dtype, - use_tma: bool, + tma: bool, + triltf32: bool, ): - if not use_tma: + if not tma: os.environ['FLA_USE_TMA'] = '0' else: os.environ['FLA_USE_TMA'] = '1' + if triltf32: + os.environ['FLA_TRIL_PRECISION'] = 'tf32x3' + else: + os.environ['FLA_TRIL_PRECISION'] = 'ieee' torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments @@ -313,6 +324,7 @@ def test_chunk_varlen( assert_close('db', ref_db, tri_db, 0.015) assert_close('dh0', ref_dh0, tri_dh0, 0.007) os.environ['FLA_USE_TMA'] = '0' + os.environ['FLA_TRIL_PRECISION'] = 'ieee' @pytest.mark.parametrize(