diff --git a/fla/ops/attn/decoding.py b/fla/ops/attn/decoding.py index 23eb5e1c9f..1f62f3a5fb 100644 --- a/fla/ops/attn/decoding.py +++ b/fla/ops/attn/decoding.py @@ -52,8 +52,8 @@ def naive_attn_decoding_kernel( i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G - bos, eos = tl.load(cu_seqlens + i_b).to(tl.int32), tl.load(cu_seqlens + i_b + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_b).to(tl.int64), tl.load(cu_seqlens + i_b + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1, ), (0, ), (BK,), (0,)) p_o = tl.make_block_ptr(o + i_bh * V, (V,), (1, ), (0, ), (BV,), (0,)) diff --git a/fla/ops/attn/parallel.py b/fla/ops/attn/parallel.py index c57810acd4..458c735beb 100644 --- a/fla/ops/attn/parallel.py +++ b/fla/ops/attn/parallel.py @@ -54,11 +54,11 @@ def parallel_attn_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T RCP_LN2: tl.constexpr = 1.4426950216 p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) @@ -226,11 +226,11 @@ def parallel_attn_bwd_kernel_dq( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T # NOTE: we must multiply RCP_LN2 after tl.dot for high precision RCP_LN2: tl.constexpr = 1.4426950216 @@ -372,11 +372,11 @@ def parallel_attn_bwd_kernel_dkv( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T RCP_LN2: tl.constexpr = 1.4426950216 p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) diff --git a/fla/ops/comba/fused_recurrent.py b/fla/ops/comba/fused_recurrent.py index 04411b8b0f..abb545409f 100644 --- a/fla/ops/comba/fused_recurrent.py +++ b/fla/ops/comba/fused_recurrent.py @@ -52,9 +52,9 @@ def fused_recurrent_comba_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/comba/utils.py b/fla/ops/comba/utils.py index a6af1cc1f4..cd409e3b74 100644 --- a/fla/ops/comba/utils.py +++ b/fla/ops/comba/utils.py @@ -44,10 +44,10 @@ def chunk_comba_cumsum_scalar_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) @@ -120,10 +120,10 @@ def chunk_comba_cumsum_scalar_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) diff --git a/fla/ops/comba/wy_fast.py b/fla/ops/comba/wy_fast.py index 5b4efed9d3..1ff963c64e 100644 --- a/fla/ops/comba/wy_fast.py +++ b/fla/ops/comba/wy_fast.py @@ -51,10 +51,10 @@ def chunk_scaled_dot_comba_pkt_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -194,10 +194,10 @@ def prepare_wy_repr_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) p_g0 = tl.make_block_ptr(g0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) @@ -322,10 +322,10 @@ def recompute_w_u_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) diff --git a/fla/ops/common/chunk_delta_h.py b/fla/ops/common/chunk_delta_h.py index c29495a2fd..9347471692 100644 --- a/fla/ops/common/chunk_delta_h.py +++ b/fla/ops/common/chunk_delta_h.py @@ -68,12 +68,12 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // HV, i_nh % HV if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -375,12 +375,12 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // HV, i_nh % HV if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/common/chunk_h.py b/fla/ops/common/chunk_h.py index 41c12da8b0..10cc3eae20 100644 --- a/fla/ops/common/chunk_h.py +++ b/fla/ops/common/chunk_h.py @@ -64,12 +64,12 @@ def chunk_fwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) boh = i_n * NS NTS = BS // BT @@ -195,13 +195,13 @@ def chunk_bwd_kernel_dh( i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) boh = i_n * NS diff --git a/fla/ops/common/chunk_h_parallel.py b/fla/ops/common/chunk_h_parallel.py index 26c1033cc1..3ddd3e9ea7 100644 --- a/fla/ops/common/chunk_h_parallel.py +++ b/fla/ops/common/chunk_h_parallel.py @@ -74,11 +74,11 @@ def chunk_fwd_kernel_h_parallel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T NT = tl.cdiv(T, BT) i_n, i_tg = i_b, i_b * NT + i_t i_nh = i_n * H + i_h @@ -176,12 +176,12 @@ def chunk_fwd_kernel_h_reduction( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -274,11 +274,11 @@ def chunk_bwd_kernel_dh_parallel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T NT = tl.cdiv(T, BT) i_n, i_tg = i_b, i_b * NT + i_t i_nh = i_n * HQ + i_hq @@ -369,12 +369,12 @@ def chunk_bwd_kernel_dh_reduction( i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/common/chunk_h_split.py b/fla/ops/common/chunk_h_split.py index 6264d87d4a..fd99f8f1e9 100644 --- a/fla/ops/common/chunk_h_split.py +++ b/fla/ops/common/chunk_h_split.py @@ -65,13 +65,13 @@ def chunk_fwd_kernel_h_split( i_ss, i_h = i_sh // H, i_sh % H if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T i_nh = i_n * H + i_h # [BK, BV] @@ -176,12 +176,12 @@ def chunk_fwd_kernel_h_reduction( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NS = tl.cdiv(T, S) boh = i_n * NS @@ -275,13 +275,13 @@ def chunk_bwd_kernel_dh_split( i_ss, i_hq = i_sh // HQ, i_sh % HQ if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T i_nh = i_n * HQ + i_hq i_h = i_hq // NG @@ -385,12 +385,12 @@ def chunk_bwd_kernel_dh_reduction( i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NS = tl.cdiv(T, S) boh = i_n * NS diff --git a/fla/ops/common/chunk_o.py b/fla/ops/common/chunk_o.py index 7999923c94..3a8cf34938 100644 --- a/fla/ops/common/chunk_o.py +++ b/fla/ops/common/chunk_o.py @@ -63,13 +63,13 @@ def chunk_fwd_kernel_o( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h // (HV // H)) * K @@ -194,13 +194,13 @@ def chunk_bwd_kernel_dqkwg( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation v += (bos * HV + i_h) * V @@ -385,13 +385,13 @@ def chunk_bwd_kernel_dv( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T b_dv = tl.zeros([BT, BV], dtype=tl.float32) @@ -487,10 +487,10 @@ def chunk_bwd_kernel_dv_local( i_b, i_h = i_bh // HV, i_bh % HV if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h // (HV // H)) * K diff --git a/fla/ops/common/chunk_scaled_dot_kkt.py b/fla/ops/common/chunk_scaled_dot_kkt.py index 83a45f122f..1507fee003 100644 --- a/fla/ops/common/chunk_scaled_dot_kkt.py +++ b/fla/ops/common/chunk_scaled_dot_kkt.py @@ -49,10 +49,10 @@ def chunk_scaled_dot_kkt_fwd_kernel( i_b, i_h = i_bh // HV, i_bh % HV if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/common/fused_chunk.py b/fla/ops/common/fused_chunk.py index 364811c30a..11fda7bb81 100644 --- a/fla/ops/common/fused_chunk.py +++ b/fla/ops/common/fused_chunk.py @@ -72,10 +72,10 @@ def fused_chunk_fwd_kernel( all = B * T if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) o_i = tl.arange(0, BT) @@ -211,10 +211,10 @@ def fused_chunk_bwd_kernel( all = B * T if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) NV = tl.cdiv(V, BV) diff --git a/fla/ops/common/fused_recurrent.py b/fla/ops/common/fused_recurrent.py index d6f146bfb4..808ce08894 100644 --- a/fla/ops/common/fused_recurrent.py +++ b/fla/ops/common/fused_recurrent.py @@ -62,9 +62,9 @@ def fused_recurrent_fwd_kernel( all = B * T if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) @@ -184,9 +184,9 @@ def fused_recurrent_bwd_kernel( all = B * T if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NV = tl.cdiv(V, BV) o_k = i_k * BK + tl.arange(0, BK) diff --git a/fla/ops/cp/chunk_delta_h.py b/fla/ops/cp/chunk_delta_h.py index e06a111483..a228ab9952 100644 --- a/fla/ops/cp/chunk_delta_h.py +++ b/fla/ops/cp/chunk_delta_h.py @@ -385,9 +385,9 @@ def merge_fwd_bwd_kernel( return # Load offsets for this sequence - ss_start = tl.load(seq_offsets + i_seq).to(tl.int32) - ss_end = tl.load(seq_offsets + i_seq + 1).to(tl.int32) - init_base = tl.load(init_offsets + i_seq).to(tl.int32) + ss_start = tl.load(seq_offsets + i_seq).to(tl.int64) + ss_end = tl.load(seq_offsets + i_seq + 1).to(tl.int64) + init_base = tl.load(init_offsets + i_seq).to(tl.int64) num_subseqs = ss_end - ss_start stride_hm_s = HV * K * (V + K) diff --git a/fla/ops/delta_rule/fused_recurrent.py b/fla/ops/delta_rule/fused_recurrent.py index d9c018b5af..e4f71e7ce3 100644 --- a/fla/ops/delta_rule/fused_recurrent.py +++ b/fla/ops/delta_rule/fused_recurrent.py @@ -47,9 +47,9 @@ def fused_recurrent_delta_rule_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) @@ -139,9 +139,9 @@ def fused_recurrent_delta_rule_bwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T mask_k = i_k * BK + tl.arange(0, BK) < K diff --git a/fla/ops/delta_rule/wy_fast.py b/fla/ops/delta_rule/wy_fast.py index f75cf65c6e..36e873f413 100644 --- a/fla/ops/delta_rule/wy_fast.py +++ b/fla/ops/delta_rule/wy_fast.py @@ -52,10 +52,10 @@ def recompute_w_u_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) @@ -117,10 +117,10 @@ def prepare_wy_repr_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) diff --git a/fla/ops/gated_delta_product/chunk_deltaproduct_h.py b/fla/ops/gated_delta_product/chunk_deltaproduct_h.py index 983440f42f..941a17f86a 100644 --- a/fla/ops/gated_delta_product/chunk_deltaproduct_h.py +++ b/fla/ops/gated_delta_product/chunk_deltaproduct_h.py @@ -62,12 +62,12 @@ def chunk_gated_delta_product_fwd_kernel_h_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * tl.cdiv(T // num_householder, BT) @@ -242,12 +242,12 @@ def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/gated_delta_product/chunk_deltaproduct_o.py b/fla/ops/gated_delta_product/chunk_deltaproduct_o.py index 7895a830ed..ce1755392b 100644 --- a/fla/ops/gated_delta_product/chunk_deltaproduct_o.py +++ b/fla/ops/gated_delta_product/chunk_deltaproduct_o.py @@ -60,13 +60,13 @@ def chunk_fwd_kernel_o( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h) * K diff --git a/fla/ops/gated_delta_rule/chunk_fwd.py b/fla/ops/gated_delta_rule/chunk_fwd.py index a0f8fe3a2f..8f87f4c660 100644 --- a/fla/ops/gated_delta_rule/chunk_fwd.py +++ b/fla/ops/gated_delta_rule/chunk_fwd.py @@ -70,10 +70,10 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT >= T: return diff --git a/fla/ops/gated_delta_rule/fused_recurrent.py b/fla/ops/gated_delta_rule/fused_recurrent.py index 9eb9fe01e8..626c579211 100644 --- a/fla/ops/gated_delta_rule/fused_recurrent.py +++ b/fla/ops/gated_delta_rule/fused_recurrent.py @@ -59,9 +59,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_k = tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/gated_delta_rule/gate.py b/fla/ops/gated_delta_rule/gate.py index ae737a6e33..f1b908acae 100644 --- a/fla/ops/gated_delta_rule/gate.py +++ b/fla/ops/gated_delta_rule/gate.py @@ -79,10 +79,10 @@ def gdn_gate_chunk_cumsum_scalar_kernel( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py index c3868000a6..fde9d1fa61 100644 --- a/fla/ops/gated_delta_rule/wy_fast.py +++ b/fla/ops/gated_delta_rule/wy_fast.py @@ -80,10 +80,10 @@ def recompute_w_u_fwd_kernel( i_b, i_h = i_bh // HV, i_bh % HV if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) @@ -160,10 +160,10 @@ def prepare_wy_repr_bwd_kernel( i_b, i_h = i_bh // HV, i_bh % HV if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) p_db = tl.make_block_ptr(db + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) diff --git a/fla/ops/gated_oja_rule/chunk_h.py b/fla/ops/gated_oja_rule/chunk_h.py index af26d8d86c..e25e78cc18 100644 --- a/fla/ops/gated_oja_rule/chunk_h.py +++ b/fla/ops/gated_oja_rule/chunk_h.py @@ -62,12 +62,12 @@ def chunk_oja_fwd_kernel_h_blockdim64( i_k, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -290,12 +290,12 @@ def chunk_oja_bwd_kernel_dhu_blockdim64( i_k, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -561,14 +561,14 @@ def chunk_gsa_bwd_k_kernel_dqkvg( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T o_i = tl.arange(0, BT) @@ -691,13 +691,13 @@ def chunk_oja_bwd_kernel_dvwg_h( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation k += (bos * H + i_h) * K diff --git a/fla/ops/gated_oja_rule/chunk_kkt.py b/fla/ops/gated_oja_rule/chunk_kkt.py index 7c482cce1e..0d4999cef9 100644 --- a/fla/ops/gated_oja_rule/chunk_kkt.py +++ b/fla/ops/gated_oja_rule/chunk_kkt.py @@ -46,10 +46,10 @@ def chunk_scaled_dot_kkt_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -109,10 +109,10 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -183,10 +183,10 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -255,10 +255,10 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( all = B * T if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) if i_t * BT + i_i * BC >= T: return diff --git a/fla/ops/gated_oja_rule/chunk_o.py b/fla/ops/gated_oja_rule/chunk_o.py index b522ffe85c..7bbbbb04e6 100644 --- a/fla/ops/gated_oja_rule/chunk_o.py +++ b/fla/ops/gated_oja_rule/chunk_o.py @@ -58,13 +58,13 @@ def chunk_oja_fwd_inter( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] @@ -129,10 +129,10 @@ def chunk_oja_fwd_intra( i_t, i_i = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_v = i_v * BV + tl.arange(0, BV) m_v = o_v < V @@ -283,11 +283,11 @@ def chunk_oja_bwd_kernel_dA( i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T o_v = i_v * BV + tl.arange(0, BV) @@ -430,14 +430,14 @@ def chunk_oja_bwd_kernel_dqk( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T o_i = tl.arange(0, BT) @@ -565,10 +565,10 @@ def chunk_oja_bwd_kernel_dv_o( i_t, i_i = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_v = i_v * BV + tl.arange(0, BV) m_v = o_v < V diff --git a/fla/ops/gated_oja_rule/fused_recurrent.py b/fla/ops/gated_oja_rule/fused_recurrent.py index e82e9bc4ac..57aa353973 100644 --- a/fla/ops/gated_oja_rule/fused_recurrent.py +++ b/fla/ops/gated_oja_rule/fused_recurrent.py @@ -52,9 +52,9 @@ def fused_recurrent_oja_fwd_kernel( i_h = i_hv // (HV // H) if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_k = tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/gated_oja_rule/wy_fast.py b/fla/ops/gated_oja_rule/wy_fast.py index e1b2135317..2c380f07bc 100644 --- a/fla/ops/gated_oja_rule/wy_fast.py +++ b/fla/ops/gated_oja_rule/wy_fast.py @@ -52,10 +52,10 @@ def recompute_w_u_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) @@ -134,10 +134,10 @@ def prepare_wy_repr_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py index a051ff50b3..3eaae514ca 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py @@ -67,10 +67,10 @@ def chunk_dplr_bwd_kernel_intra( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT >= T: return @@ -273,10 +273,10 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T_len = eos - bos else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T T_len = T if i_t * BT >= T_len: @@ -424,13 +424,13 @@ def chunk_dplr_bwd_dgk_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = (i_b * NT + i_t).to(tl.int32) - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T stride_qk = H * K dgk += (bos * H + i_h) * K diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py index 2a3a475fca..d6868761e5 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py @@ -61,10 +61,10 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT >= T: return @@ -188,10 +188,10 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T_len = eos - bos else: - bos = i_b * T + bos = tl.cast(i_b, tl.int64) * T T_len = T if i_t * BT >= T_len: diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py index 98e25b9b26..4fc9bb1cb3 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py @@ -60,12 +60,12 @@ def chunk_dplr_bwd_kernel_dhu( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py index 5e29797524..ade7f785e4 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py @@ -60,12 +60,12 @@ def chunk_dplr_fwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT o_k = i_k * BK + tl.arange(0, BK) diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py index e7937064c3..f61705dc18 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py @@ -54,10 +54,10 @@ def chunk_dplr_bwd_kernel_dAu( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) b_dA_qk = tl.zeros([BT, BT], dtype=tl.float32) b_dA_qb = tl.zeros([BT, BT], dtype=tl.float32) @@ -138,13 +138,13 @@ def chunk_dplr_bwd_o_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation v += (bos * H + i_h) * V @@ -257,13 +257,13 @@ def chunk_dplr_bwd_kernel_dv( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T b_dv = tl.zeros([BT, BV], dtype=tl.float32) diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py index 53b6e2bf74..47d003baf4 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py @@ -58,13 +58,13 @@ def chunk_dplr_fwd_kernel_o( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T b_o = tl.zeros([BT, BV], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): diff --git a/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py index ed621379b6..98a9f1ce7e 100644 --- a/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py +++ b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py @@ -59,9 +59,9 @@ def fused_recurrent_dplr_delta_rule_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_k = tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py b/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py index dc4d3c1df8..3bba0844c3 100644 --- a/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py @@ -57,10 +57,10 @@ def prepare_wy_repr_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_Aak_t = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) p_Aab_inv_t = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) diff --git a/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py b/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py index f0cdc5fa04..b87a77c105 100644 --- a/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py @@ -42,10 +42,10 @@ def prepare_wy_repr_fwd_kernel_chunk32( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_Aab = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) p_Aab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) b_A_ab = tl.load(p_Aab, boundary_check=(0, 1)) @@ -89,10 +89,10 @@ def prepare_wy_repr_fwd_kernel_chunk64( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_A1 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) p_A2 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) @@ -175,10 +175,10 @@ def wu_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_s = tl.arange(0, BT) p_A_ab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) diff --git a/fla/ops/generalized_delta_rule/iplr/chunk.py b/fla/ops/generalized_delta_rule/iplr/chunk.py index 1c8c8bd803..f891bf6a2f 100644 --- a/fla/ops/generalized_delta_rule/iplr/chunk.py +++ b/fla/ops/generalized_delta_rule/iplr/chunk.py @@ -67,12 +67,12 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -151,13 +151,13 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_o( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h) * K diff --git a/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py index a9ec673241..4be4bcb1a4 100644 --- a/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py +++ b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py @@ -55,9 +55,9 @@ def fused_recurrent_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T p_q = q + (bos * H + i_h) * K + tl.arange(0, BK) p_k = k + (bos * H + i_h) * K + tl.arange(0, BK) @@ -160,9 +160,9 @@ def fused_recurrent_bwd_kernel( da += i_v * B * H * K * T if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T mask_k = tl.arange(0, BK) < K mask_v = (tl.arange(0, BV) + i_v * BV) < V diff --git a/fla/ops/generalized_delta_rule/iplr/wy_fast.py b/fla/ops/generalized_delta_rule/iplr/wy_fast.py index 844d7615a8..fa4bd54693 100644 --- a/fla/ops/generalized_delta_rule/iplr/wy_fast.py +++ b/fla/ops/generalized_delta_rule/iplr/wy_fast.py @@ -45,10 +45,10 @@ def prepare_wy_repr_fwd_kernel_chunk32( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): @@ -100,10 +100,10 @@ def prepare_wy_repr_fwd_kernel_chunk64( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T b_A = tl.zeros([BC, BC], dtype=tl.float32) b_A2 = tl.zeros([BC, BC], dtype=tl.float32) @@ -185,10 +185,10 @@ def wu_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index c2a1c82ad1..4739a54b66 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -55,10 +55,10 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -128,10 +128,10 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( i_j = i_i if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -205,11 +205,11 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( i_j = i_i if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T if i_t * BT + i_i * BC >= T: @@ -277,11 +277,11 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T if i_t * BT + i_c * BC >= T: @@ -337,7 +337,7 @@ def chunk_gla_fwd_kernel_o( i_tg = i_t.to(tl.int64) i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) @@ -419,10 +419,10 @@ def chunk_gla_bwd_kernel_intra( i_k, i_i = i_kc // NC, i_kc % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) if i_t * BT + i_i * BC >= T: return @@ -550,10 +550,10 @@ def chunk_gla_bwd_kernel_dA( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): @@ -608,13 +608,13 @@ def chunk_gla_bwd_kernel_dv( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) p_do = tl.make_block_ptr(do + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) @@ -692,13 +692,13 @@ def chunk_gla_bwd_kernel_inter( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K diff --git a/fla/ops/gsa/chunk.py b/fla/ops/gsa/chunk.py index 6ec197d2ef..44366656e9 100644 --- a/fla/ops/gsa/chunk.py +++ b/fla/ops/gsa/chunk.py @@ -63,13 +63,13 @@ def chunk_gsa_fwd_k_kernel_inter( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] @@ -134,10 +134,10 @@ def chunk_gsa_fwd_k_kernel_intra( i_t, i_i = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, BC) o_v = i_v * BV + tl.arange(0, BV) @@ -224,11 +224,11 @@ def chunk_gsa_bwd_k_kernel_dA( i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T o_v = i_v * BV + tl.arange(0, BV) @@ -335,14 +335,14 @@ def chunk_gsa_bwd_k_kernel_dqkvg( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T o_i = tl.arange(0, BT) @@ -452,10 +452,10 @@ def chunk_gsa_bwd_k_kernel_intra_dvg( i_t, i_i = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, BC) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/hgrn/chunk.py b/fla/ops/hgrn/chunk.py index b39b4c6eac..bf6a59b047 100644 --- a/fla/ops/hgrn/chunk.py +++ b/fla/ops/hgrn/chunk.py @@ -66,10 +66,10 @@ def chunk_hgrn_fwd_kernel_h( o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D - p_x = x + i_b * T * D + i_t * BT * D + o_d - p_g = g + i_b * T * D + i_t * BT * D + o_d - p_gc = gc + i_b * T * D + i_t * BT * D + o_d - p_o = o + i_b * T * D + i_t * BT * D + o_d + p_x = x + tl.cast(i_b, tl.int64) * T * D + i_t * BT * D + o_d + p_g = g + tl.cast(i_b, tl.int64) * T * D + i_t * BT * D + o_d + p_gc = gc + tl.cast(i_b, tl.int64) * T * D + i_t * BT * D + o_d + p_o = o + tl.cast(i_b, tl.int64) * T * D + i_t * BT * D + o_d b_h = tl.zeros([BD], dtype=tl.float32) b_gc = tl.zeros([BD], dtype=tl.float32) @@ -108,11 +108,11 @@ def chunk_hgrn_fwd_kernel_o( mask = o_d < D for i_t in range(1, tl.cdiv(T, BT)): - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_gc = tl.make_block_ptr(gc + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_o = tl.make_block_ptr(o + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) # [BD,] - b_h0 = tl.load(o + i_b * T * D + i_t * BT * D - D + o_d, mask=mask, other=0).to(tl.float32) + b_h0 = tl.load(o + tl.cast(i_b, tl.int64) * T * D + i_t * BT * D - D + o_d, mask=mask, other=0).to(tl.float32) # [BT, BD] b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) @@ -146,15 +146,15 @@ def chunk_hgrn_bwd_kernel_h( BC = min(BT, T - i_t * BT) NT = tl.num_programs(1) - p_g = g + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_gc = gc + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_dx = dx + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_do = do + (i_b * T + i_t * BT + BC - 1) * D + o_d + p_g = g + (tl.cast(i_b, tl.int64) * T + i_t * BT + BC - 1) * D + o_d + p_gc = gc + (tl.cast(i_b, tl.int64) * T + i_t * BT + BC - 1) * D + o_d + p_dx = dx + (tl.cast(i_b, tl.int64) * T + i_t * BT + BC - 1) * D + o_d + p_do = do + (tl.cast(i_b, tl.int64) * T + i_t * BT + BC - 1) * D + o_d if i_t == NT - 1: b_gc = tl.zeros([BD], dtype=tl.float32) else: - b_gc = tl.load(g + (i_b * T + i_t * BT + BT) * D + o_d, mask=mask, other=0).to(tl.float32) + b_gc = tl.load(g + (tl.cast(i_b, tl.int64) * T + i_t * BT + BT) * D + o_d, mask=mask, other=0).to(tl.float32) b_dh = tl.zeros([BD], dtype=tl.float32) for _ in range(BC - 1, -1, -1): tl.store(p_gc, b_gc.to(p_gc.dtype.element_ty), mask=mask) @@ -195,15 +195,16 @@ def chunk_hgrn_bwd_kernel_o( mask = o_d < D for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): - p_g = tl.make_block_ptr(g + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT - 1, i_d * BD), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_dg = tl.make_block_ptr(dg + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_g = tl.make_block_ptr(g + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_gc = tl.make_block_ptr(gc + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_o = tl.make_block_ptr(o + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), + (i_t * BT - 1, i_d * BD), (BT, BD), (1, 0)) + p_dx = tl.make_block_ptr(dx + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_dg = tl.make_block_ptr(dg + tl.cast(i_b, tl.int64) * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) # [BD,] mask_t = mask & ((i_t + 1) * BT < T) - b_ht = tl.load(dx + i_b * T * D + (i_t + 1) * BT * D + o_d, mask=mask_t, other=0).to(tl.float32) + b_ht = tl.load(dx + tl.cast(i_b, tl.int64) * T * D + (i_t + 1) * BT * D + o_d, mask=mask_t, other=0).to(tl.float32) # [BT, BD] b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) diff --git a/fla/ops/hgrn/fused_recurrent.py b/fla/ops/hgrn/fused_recurrent.py index 0c38f56c29..02ac546add 100644 --- a/fla/ops/hgrn/fused_recurrent.py +++ b/fla/ops/hgrn/fused_recurrent.py @@ -45,9 +45,9 @@ def fused_recurrent_hgrn_fwd_kernel( i_d, i_n = tl.program_id(0), tl.program_id(1) if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D @@ -110,9 +110,9 @@ def fused_recurrent_hgrn_bwd_kernel( i_d, i_n = tl.program_id(0), tl.program_id(1) if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D diff --git a/fla/ops/kda/chunk_bwd.py b/fla/ops/kda/chunk_bwd.py index b46becf38c..af0f726e70 100644 --- a/fla/ops/kda/chunk_bwd.py +++ b/fla/ops/kda/chunk_bwd.py @@ -69,10 +69,10 @@ def chunk_kda_bwd_kernel_dAv( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h) * K diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py index 9a1e21aabc..70be49bec2 100644 --- a/fla/ops/kda/chunk_intra.py +++ b/fla/ops/kda/chunk_intra.py @@ -75,10 +75,10 @@ def chunk_kda_fwd_kernel_inter_solve_fused( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT >= T: return @@ -397,10 +397,10 @@ def chunk_kda_bwd_kernel_intra( all = B * T if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) i_ti = i_t * BT + i_i * BC if i_ti >= T: @@ -662,10 +662,10 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T i_ti = i_t * BT + i_i * BC if i_ti >= T: diff --git a/fla/ops/kda/chunk_intra_token_parallel.py b/fla/ops/kda/chunk_intra_token_parallel.py index 487e9aac18..534b57ea7e 100644 --- a/fla/ops/kda/chunk_intra_token_parallel.py +++ b/fla/ops/kda/chunk_intra_token_parallel.py @@ -58,14 +58,14 @@ def chunk_kda_fwd_kernel_intra_token_parallel( for _ in range(20): if left < right: mid = (left + right) // 2 - if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int64): right = mid else: left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) i_t = i_tg - bos else: bos = (i_tg // T) * T diff --git a/fla/ops/kda/fused_recurrent.py b/fla/ops/kda/fused_recurrent.py index 675dbc25e2..17b699b573 100644 --- a/fla/ops/kda/fused_recurrent.py +++ b/fla/ops/kda/fused_recurrent.py @@ -85,9 +85,9 @@ def fused_recurrent_kda_fwd_kernel( tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T if T == 0: # no tokens to process for this sequence diff --git a/fla/ops/kda/gate.py b/fla/ops/kda/gate.py index f1ac2da6f8..b0fbd62447 100644 --- a/fla/ops/kda/gate.py +++ b/fla/ops/kda/gate.py @@ -383,10 +383,10 @@ def kda_gate_chunk_cumsum_vector_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) diff --git a/fla/ops/kda/wy_fast.py b/fla/ops/kda/wy_fast.py index 7c84ed97cf..9d171ee883 100644 --- a/fla/ops/kda/wy_fast.py +++ b/fla/ops/kda/wy_fast.py @@ -57,10 +57,10 @@ def recompute_w_u_fwd_kda_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) @@ -146,10 +146,10 @@ def prepare_wy_repr_bwd_kda_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) diff --git a/fla/ops/log_linear_attn/chunk.py b/fla/ops/log_linear_attn/chunk.py index 87b82c7263..5e91942ec1 100644 --- a/fla/ops/log_linear_attn/chunk.py +++ b/fla/ops/log_linear_attn/chunk.py @@ -76,12 +76,12 @@ def chunkwise_fwd_kernel( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_i = tl.arange(0, BT) @@ -707,12 +707,12 @@ def copy_input_kernel( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T offset = tl.load(offsets + i_n) input_offset = -1 * (offset % BT) @@ -853,12 +853,12 @@ def copy_last_chunk_kernel( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T seq_offset = (T // BT) * BT @@ -955,12 +955,12 @@ def chunkwise_bwd_kernel_dhg( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T b_dh = tl.zeros([BK, V], dtype=tl.float32) @@ -1073,12 +1073,12 @@ def chunkwise_bwd_kernel_hdqgl( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T b_h = tl.zeros([V, K], dtype=tl.float32) @@ -1216,12 +1216,12 @@ def chunkwise_bwd_kernel_dkg( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_i = tl.arange(0, BT) o_t = i_t * BT + o_i @@ -1303,12 +1303,12 @@ def chunkwise_bwd_kernel_dv( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -1379,12 +1379,12 @@ def chunkwise_bwd_kernel_diag( if IS_VARLEN: bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_i = tl.arange(0, BT) i_idx = o_i[:, None] # BT x 1 diff --git a/fla/ops/mesa_net/chunk_cg_solver_bwd.py b/fla/ops/mesa_net/chunk_cg_solver_bwd.py index 9eef3c4c28..dc9673e37e 100644 --- a/fla/ops/mesa_net/chunk_cg_solver_bwd.py +++ b/fla/ops/mesa_net/chunk_cg_solver_bwd.py @@ -58,13 +58,13 @@ def chunk_fwd_mesa_cg_dim64_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/mesa_net/chunk_cg_solver_fwd.py b/fla/ops/mesa_net/chunk_cg_solver_fwd.py index 9ac727b257..38ffb33da3 100644 --- a/fla/ops/mesa_net/chunk_cg_solver_fwd.py +++ b/fla/ops/mesa_net/chunk_cg_solver_fwd.py @@ -61,13 +61,13 @@ def chunk_fwd_mesa_cg_dim64_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/mesa_net/chunk_h_fwd.py b/fla/ops/mesa_net/chunk_h_fwd.py index 60f03f3b71..591f5f9d5e 100644 --- a/fla/ops/mesa_net/chunk_h_fwd.py +++ b/fla/ops/mesa_net/chunk_h_fwd.py @@ -57,13 +57,13 @@ def chunk_mesa_net_fwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) boh = i_n * NS diff --git a/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py b/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py index 47a2dc4157..2aac261195 100644 --- a/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py +++ b/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py @@ -48,13 +48,13 @@ def chunk_mesa_net_h_kk_bwd_intra_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py b/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py index 47f108c01e..65ed48057c 100644 --- a/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py +++ b/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py @@ -60,13 +60,13 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py b/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py index 5b92278957..b9b192700f 100644 --- a/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py +++ b/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py @@ -58,13 +58,13 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -176,13 +176,13 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/nsa/compression.py b/fla/ops/nsa/compression.py index 85896a9ec7..b53ae5c542 100644 --- a/fla/ops/nsa/compression.py +++ b/fla/ops/nsa/compression.py @@ -54,11 +54,11 @@ def parallel_nsa_compression_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T boc = i_b * tl.cdiv(T, BS) p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) @@ -161,11 +161,11 @@ def parallel_nsa_compression_bwd_kernel_dq( all = B * T if IS_VARLEN: i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T boc = i_b * tl.cdiv(T, BS) q += (bos + i_t) * HQ*K @@ -268,13 +268,13 @@ def parallel_nsa_compression_bwd_kernel_dkv( if IS_VARLEN: i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int32), tl.load(chunk_indices + i_c * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) # the number of compression representations in total TC = tl.cdiv(T, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T boc = i_b * tl.cdiv(T, BS) p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) diff --git a/fla/ops/nsa/parallel.py b/fla/ops/nsa/parallel.py index 31b55f3e54..0c0802b0ab 100644 --- a/fla/ops/nsa/parallel.py +++ b/fla/ops/nsa/parallel.py @@ -66,11 +66,11 @@ def parallel_nsa_kernel_topk( if IS_VARLEN: i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T boc = i_b * tl.cdiv(T, BS) p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) @@ -208,10 +208,10 @@ def parallel_nsa_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T k += (bos * H + i_h) * K v += (bos * H + i_h) * V @@ -283,14 +283,15 @@ def parallel_nsa_kernel_mask( i_t, i_b, i_hs = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_h, i_s = i_hs // S, i_hs % S - b_i = tl.load(block_indices + i_b * T * H * S + i_t * H * S + i_h * S + i_s) + b_i = tl.load(block_indices + tl.cast(i_b, tl.int64) * T * H * S + i_t * H * S + i_h * S + i_s) if USE_BLOCK_COUNTS: - b_m = b_i * BS <= i_t and i_s < tl.load(block_counts + i_b * T * H + i_t * H + i_h) + b_m = b_i * BS <= i_t and i_s < tl.load(block_counts + tl.cast(i_b, tl.int64) * T * H + i_t * H + i_h) else: b_m = b_i * BS <= i_t if b_i < NS and b_i >= 0: - tl.store(block_mask + i_b * T * H * NS + i_t * H * NS + i_h * NS + b_i, b_m.to(block_mask.dtype.element_ty)) + tl.store(block_mask + tl.cast(i_b, tl.int64) * T * H * NS + i_t * + H * NS + i_h * NS + b_i, b_m.to(block_mask.dtype.element_ty)) @triton.heuristics({ @@ -339,10 +340,10 @@ def parallel_nsa_bwd_kernel_dq( all = B * T if IS_VARLEN: i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T q += (bos + i_t) * HQ*K do += (bos + i_t) * HQ*V @@ -447,10 +448,10 @@ def parallel_nsa_bwd_kernel_dkv( all = B * T if IS_VARLEN: i_n, i_s = tl.load(chunk_indices + i_s * 2).to(tl.int32), tl.load(chunk_indices + i_s * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s * BS, 0), (BS, BK), (1, 0)) p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s * BS, i_v * BV), (BS, BV), (1, 0)) diff --git a/fla/ops/path_attn/cumprod_householder_bwd.py b/fla/ops/path_attn/cumprod_householder_bwd.py index c177a192b3..eba7575892 100644 --- a/fla/ops/path_attn/cumprod_householder_bwd.py +++ b/fla/ops/path_attn/cumprod_householder_bwd.py @@ -36,15 +36,15 @@ def chunk_cumprod_householder_bwd_kernel( if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T boh = i_n * tl.cdiv(T, BT) boh_large = i_n * tl.cdiv(T, S) diff --git a/fla/ops/path_attn/cumprod_householder_fwd.py b/fla/ops/path_attn/cumprod_householder_fwd.py index f23bf167f1..cb622e7a1c 100644 --- a/fla/ops/path_attn/cumprod_householder_fwd.py +++ b/fla/ops/path_attn/cumprod_householder_fwd.py @@ -40,16 +40,16 @@ def chunk_cumprod_householder_fwd_kernel( if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T boh = i_n * tl.cdiv(T, BT) boh_large = i_n * tl.cdiv(T, S) diff --git a/fla/ops/path_attn/intra_chunk_preprocess_bwd.py b/fla/ops/path_attn/intra_chunk_preprocess_bwd.py index 11e7047bdd..ca4ca0511c 100644 --- a/fla/ops/path_attn/intra_chunk_preprocess_bwd.py +++ b/fla/ops/path_attn/intra_chunk_preprocess_bwd.py @@ -33,10 +33,10 @@ def intra_chunk_preprocess_bwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dw_beta = tl.zeros([BT, BK], dtype=tl.float32) diff --git a/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py b/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py index a1ed94b8e0..6f356af6c6 100644 --- a/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py +++ b/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py @@ -57,11 +57,11 @@ def chunk_transform_qk_bwd_kernel_prepare( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) - T = eos - bos - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/path_attn/intra_chunk_preprocess_fwd.py b/fla/ops/path_attn/intra_chunk_preprocess_fwd.py index e28671771e..dc3a17ba86 100644 --- a/fla/ops/path_attn/intra_chunk_preprocess_fwd.py +++ b/fla/ops/path_attn/intra_chunk_preprocess_fwd.py @@ -52,10 +52,10 @@ def intra_chunk_preprocess_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T sm_scale = scale * 1.44269504 # offset calculations diff --git a/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py b/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py index 1a63df91a2..c3f602f905 100644 --- a/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py +++ b/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py @@ -54,12 +54,12 @@ def parallel_path_bwd_dkv_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + boh_large = tl.load(split_offsets + i_n).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T boh_large = i_n * tl.cdiv(T, S) # offset calculations diff --git a/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py b/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py index 8cefbf88c5..c096359140 100644 --- a/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py +++ b/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py @@ -56,11 +56,11 @@ def parallel_path_bwd_dq_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + boh_large = tl.load(split_offsets + i_n).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T boh_large = i_n * tl.cdiv(T, S) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T diff --git a/fla/ops/path_attn/parallel_path_bwd_intra.py b/fla/ops/path_attn/parallel_path_bwd_intra.py index d595fcd6e9..ae9bf192c7 100644 --- a/fla/ops/path_attn/parallel_path_bwd_intra.py +++ b/fla/ops/path_attn/parallel_path_bwd_intra.py @@ -34,11 +34,11 @@ def parallel_path_bwd_intra_chunk_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T # offset calculations k += (bos * H + i_h) * K # GQA when H!=HQ diff --git a/fla/ops/path_attn/parallel_path_fwd.py b/fla/ops/path_attn/parallel_path_fwd.py index 081c70cb18..6929c2a7d4 100644 --- a/fla/ops/path_attn/parallel_path_fwd.py +++ b/fla/ops/path_attn/parallel_path_fwd.py @@ -51,11 +51,11 @@ def parallel_path_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) b_q = tl.zeros([BT, BK], dtype=tl.float32) diff --git a/fla/ops/path_attn/prepare_k_cache.py b/fla/ops/path_attn/prepare_k_cache.py index 1ed8ee72b6..4105d24902 100644 --- a/fla/ops/path_attn/prepare_k_cache.py +++ b/fla/ops/path_attn/prepare_k_cache.py @@ -30,11 +30,11 @@ def parallel_path_fwd_kernel_prepare_k_cache( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T k += (bos * H + i_h) * K k_new += (bos * H + i_h) * K diff --git a/fla/ops/path_attn/transform_q.py b/fla/ops/path_attn/transform_q.py index 011792bf6d..5e4d7f37bf 100644 --- a/fla/ops/path_attn/transform_q.py +++ b/fla/ops/path_attn/transform_q.py @@ -41,11 +41,11 @@ def transform_q_fwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: i_n = i_b - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T # boh = i_n * tl.cdiv(T, BS) p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) b_q = tl.zeros([BT, BK], dtype=tl.float32) diff --git a/fla/ops/rwkv6/chunk.py b/fla/ops/rwkv6/chunk.py index 4dd8983f39..990a6f674e 100644 --- a/fla/ops/rwkv6/chunk.py +++ b/fla/ops/rwkv6/chunk.py @@ -60,10 +60,10 @@ def chunk_rwkv6_fwd_cumsum_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, BT) m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.).to(tl.float32) @@ -147,10 +147,10 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -223,10 +223,10 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( i_j = i_i if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if i_t * BT + i_i * BC >= T: return @@ -301,11 +301,11 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( i_j = i_i if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T if i_t * BT + i_i * BC >= T: @@ -373,11 +373,11 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T all = B * T if i_t * BT + i_c * BC >= T: @@ -437,12 +437,12 @@ def chunk_rwkv6_bwd_kernel_dh( i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -514,10 +514,10 @@ def chunk_rwkv6_bwd_kernel_intra( i_t, i_i = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T + T = (eos - bos).to(tl.int32) if i_t * BT + i_i * BC >= T: return @@ -666,13 +666,13 @@ def chunk_rwkv6_bwd_kernel_inter( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K diff --git a/fla/ops/rwkv6/fused_recurrent.py b/fla/ops/rwkv6/fused_recurrent.py index cf56b2cd3e..55d287096b 100644 --- a/fla/ops/rwkv6/fused_recurrent.py +++ b/fla/ops/rwkv6/fused_recurrent.py @@ -57,9 +57,9 @@ def fused_recurrent_rwkv6_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T o_k = i_k * BK + tl.arange(0, BK) @@ -143,9 +143,9 @@ def fused_recurrent_rwkv6_bwd_kernel_dq( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T o_k = i_k * BK + tl.arange(0, BK) @@ -235,9 +235,9 @@ def fused_recurrent_rwkv6_bwd_kernel_dkv( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T all = B * T o_k = i_k * BK + tl.arange(0, BK) @@ -323,10 +323,10 @@ def fused_recurrent_rwkv6_bwd_kernel_dw( i_k, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) o_i = tl.arange(0, BT) diff --git a/fla/ops/rwkv7/fused_addcmul.py b/fla/ops/rwkv7/fused_addcmul.py index 59d91f812b..e27c2c6b3c 100644 --- a/fla/ops/rwkv7/fused_addcmul.py +++ b/fla/ops/rwkv7/fused_addcmul.py @@ -65,7 +65,7 @@ def fused_addcmul_fwd_kernel( ): i_b, i_t = tl.program_id(0), tl.program_id(1) * BT - bos = i_b * (T + T_OFFSET) + bos = tl.cast(i_b, tl.int64) * (T + T_OFFSET) t_vec = i_t + T_OFFSET + tl.arange(0, BT) mask_t = t_vec < (T + T_OFFSET) o_d = tl.arange(0, BD)[None, :] @@ -143,7 +143,7 @@ def addcmul_bwd_kernel1( mask_d = d_idx < D mask = mask_t & mask_d - offset_base = i_b * (T + T_OFFSET) * D + offset_base = tl.cast(i_b, tl.int64) * (T + T_OFFSET) * D x_idx = (offset_base + t_idx * D + d_idx).to(tl.uint32) b_dxr = tl.load(dxr + x_idx, mask=mask).to(DTYPE) diff --git a/fla/ops/rwkv7/fused_k_update.py b/fla/ops/rwkv7/fused_k_update.py index b5410b0af0..96f2c6a264 100644 --- a/fla/ops/rwkv7/fused_k_update.py +++ b/fla/ops/rwkv7/fused_k_update.py @@ -40,15 +40,15 @@ def k_update_fwd_kernel_short( i_b, i_t = tl.program_id(0), tl.program_id(1) if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_b).to(tl.int64) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int64) g_t = bos + i_t if g_t >= eos: return offset = g_t * D else: g_t = i_t - offset = i_b * T * D + g_t * D + offset = tl.cast(i_b, tl.int64) * T * D + g_t * D o_d = tl.arange(0, BD) m_d = o_d < D @@ -85,12 +85,12 @@ def k_update_fwd_kernel_long( if IS_VARLEN: i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, eos - bos) else: - bos = i_b * T + bos = tl.cast(i_b, tl.int64) * T eos = (i_b + 1) * T t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, T) @@ -132,11 +132,11 @@ def k_update_bwd_kernel_short( i_b, i_t_base = tl.program_id(0), tl.program_id(1) * BT if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_b).to(tl.int64) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int64) seq_len = eos - bos else: - bos = i_b * T + bos = tl.cast(i_b, tl.int64) * T eos = (i_b + 1) * T seq_len = T @@ -185,12 +185,12 @@ def k_update_bwd_kernel_long( if IS_VARLEN: i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, eos - bos) else: - bos = i_b * T + bos = tl.cast(i_b, tl.int64) * T eos = (i_b + 1) * T t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, T) diff --git a/fla/ops/rwkv7/fused_recurrent.py b/fla/ops/rwkv7/fused_recurrent.py index dae2216405..ac7b56dcc5 100644 --- a/fla/ops/rwkv7/fused_recurrent.py +++ b/fla/ops/rwkv7/fused_recurrent.py @@ -63,9 +63,9 @@ def fused_recurrent_rwkv7_fwd_kernel( if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos + T = (eos - bos).to(tl.int32) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T o_k = tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) diff --git a/fla/ops/simple_gla/parallel.py b/fla/ops/simple_gla/parallel.py index 64ba10443b..e4bc532bf8 100644 --- a/fla/ops/simple_gla/parallel.py +++ b/fla/ops/simple_gla/parallel.py @@ -76,10 +76,10 @@ def parallel_simple_gla_fwd_kernel( all = B * T if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T q += (bos * H + i_h) * K k += (bos * H + i_h) * K @@ -417,10 +417,10 @@ def parallel_simple_gla_bwd_kernel( if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T q += (bos * H + i_h) * K k += (bos * H + i_h) * K diff --git a/fla/ops/ttt/chunk.py b/fla/ops/ttt/chunk.py index 82893225ae..078b4cf078 100755 --- a/fla/ops/ttt/chunk.py +++ b/fla/ops/ttt/chunk.py @@ -63,12 +63,12 @@ def chunk_ttt_linear_fwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -163,13 +163,13 @@ def chunk_ttt_linear_fwd_kernel_o( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h) * K @@ -261,12 +261,12 @@ def chunk_ttt_linear_bwd_kernel_h( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -354,10 +354,10 @@ def chunk_ttt_linear_bwd_kernel_dv_local( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation q += (bos * H + i_h) * K @@ -450,12 +450,12 @@ def chunk_ttt_linear_bwd_kernel_norm( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -595,13 +595,13 @@ def chunk_bwd_kernel_dqke( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T # offset calculation v += (bos * H + i_h) * V diff --git a/fla/ops/ttt/fused_chunk.py b/fla/ops/ttt/fused_chunk.py index 4175307706..80dd4e3a19 100755 --- a/fla/ops/ttt/fused_chunk.py +++ b/fla/ops/ttt/fused_chunk.py @@ -63,11 +63,11 @@ def fused_chunk_ttt_linear_fwd_kernel( i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: - bos, eos = i_n * T, i_n * T + T + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) o_i = tl.arange(0, BT) @@ -185,7 +185,7 @@ def fused_chunk_ttt_linear_bwd_kernel_h( ): i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H - bos, _ = i_n * T, i_n * T + T + bos, _ = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT @@ -313,7 +313,7 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( ): i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H - bos, _ = i_n * T, i_n * T + T + bos, _ = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T NT = tl.cdiv(T, BT) boh = i_n * NT diff --git a/fla/ops/utils/cumsum.py b/fla/ops/utils/cumsum.py index c054774f22..95a905d7dd 100644 --- a/fla/ops/utils/cumsum.py +++ b/fla/ops/utils/cumsum.py @@ -47,10 +47,10 @@ def chunk_local_cumsum_scalar_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if HEAD_FIRST: p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) @@ -104,10 +104,10 @@ def chunk_local_cumsum_vector_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T if HEAD_FIRST: p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) @@ -158,10 +158,10 @@ def chunk_global_cumsum_scalar_kernel( i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T + T = (eos - bos).to(tl.int32) b_z = tl.zeros([], dtype=tl.float32) NT = tl.cdiv(T, BT) @@ -220,10 +220,10 @@ def chunk_global_cumsum_vector_kernel( i_s, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos + bos, eos = tl.cast(i_n, tl.int64) * T, tl.cast(i_n, tl.int64) * T + T + T = (eos - bos).to(tl.int32) b_z = tl.zeros([BS], dtype=tl.float32) NT = tl.cdiv(T, BT) diff --git a/fla/ops/utils/index.py b/fla/ops/utils/index.py index 01a08c3b7e..9fd70b6290 100644 --- a/fla/ops/utils/index.py +++ b/fla/ops/utils/index.py @@ -28,8 +28,8 @@ def prepare_position_ids_kernel( B: tl.constexpr, ): i_n = tl.program_id(0) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) o = tl.arange(0, B) for i in range(0, tl.cdiv(T, B) * B, B): diff --git a/fla/ops/utils/matmul.py b/fla/ops/utils/matmul.py index 34ed273c1a..f2e16acd31 100644 --- a/fla/ops/utils/matmul.py +++ b/fla/ops/utils/matmul.py @@ -99,7 +99,7 @@ def matmul_kernel( # `p_a` is a block of [BM, BK] pointers # `p_b` is a block of [BK, BN] pointers # See above `Pointer Arithmetic` section for details - a_batch_ptr = a + i_b * stride_ab + a_batch_ptr = a + tl.cast(i_b, tl.int64) * stride_ab o_am = (i_m * BM + tl.arange(0, BM)) % M o_bn = (i_n * BN + tl.arange(0, BN)) % N o_k = tl.arange(0, BK) @@ -148,7 +148,7 @@ def matmul_kernel( # ----------------------------------------------------------- # Write back the block of the output matrix C with masks. - c_batch_ptr = c + i_b * stride_cb + c_batch_ptr = c + tl.cast(i_b, tl.int64) * stride_cb p_c = c_batch_ptr + stride_cm * o_cm[:, None] + stride_cn * o_cn[None, :] tl.store(p_c, b_c.to(c.dtype.element_ty), mask=mask) diff --git a/fla/ops/utils/pack.py b/fla/ops/utils/pack.py index 7745691055..7438a00a95 100644 --- a/fla/ops/utils/pack.py +++ b/fla/ops/utils/pack.py @@ -38,7 +38,7 @@ def packunpack_sequence_kernel( i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1) - T = eos - bos + T = (eos - bos).to(tl.int32) if PADDING_SIDE == 'left': NP = S - T if i_s < NP: @@ -53,11 +53,11 @@ def packunpack_sequence_kernel( mask = o_d < D if PACK: - b_x = tl.load(x + (i_b * S + i_s) * D + o_d, mask=mask) + b_x = tl.load(x + (tl.cast(i_b, tl.int64) * S + i_s) * D + o_d, mask=mask) tl.store(y + i_t * D + o_d, b_x, mask=mask) else: b_x = tl.load(x + i_t * D + o_d, mask=mask) - tl.store(y + (i_b * S + i_s) * D + o_d, b_x, mask=mask) + tl.store(y + (tl.cast(i_b, tl.int64) * S + i_s) * D + o_d, b_x, mask=mask) def pack_sequence_fwdbwd( diff --git a/fla/ops/utils/pooling.py b/fla/ops/utils/pooling.py index 8c8218c1b3..22af6712e4 100644 --- a/fla/ops/utils/pooling.py +++ b/fla/ops/utils/pooling.py @@ -43,13 +43,13 @@ def mean_pooling_fwd_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) @@ -90,13 +90,13 @@ def mean_pooling_bwd_kernel( if IS_VARLEN: i_tg = i_t i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) diff --git a/fla/ops/utils/solve_tril.py b/fla/ops/utils/solve_tril.py index 902affb885..f62cfbbc22 100644 --- a/fla/ops/utils/solve_tril.py +++ b/fla/ops/utils/solve_tril.py @@ -50,10 +50,10 @@ def solve_tril_16x16_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, 16) m_A = o_i[:, None] > o_i[None, :] m_I = o_i[:, None] == o_i[None, :] @@ -118,10 +118,10 @@ def merge_16x16_to_32x32_inverse_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, 16) m_A = o_i[:, None] > o_i[None, :] @@ -207,10 +207,10 @@ def merge_16x16_to_64x64_inverse_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) else: - bos, eos = i_b * T, i_b * T + T + bos, eos = tl.cast(i_b, tl.int64) * T, tl.cast(i_b, tl.int64) * T + T o_i = tl.arange(0, 16) m_A = o_i[:, None] > o_i[None, :] diff --git a/tests/ops/test_int32_overflow.py b/tests/ops/test_int32_overflow.py new file mode 100644 index 0000000000..6d189db3a0 --- /dev/null +++ b/tests/ops/test_int32_overflow.py @@ -0,0 +1,141 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +""" +Tests for int32 overflow in Triton kernel pointer arithmetic. + +When B*T*H*K > INT32_MAX (~2.15B), tl.program_id() (int32) multiplied by +large strides can overflow, causing illegal memory accesses or wrong results. + +These tests use B=4096, T=576, H=8, K=128, V=128 which gives +B*T*H*K = 2.4B > INT32_MAX. They pass with int64 index casts and +crash (illegal memory access) without them. + +Requires ~20GB GPU memory. Run with CUDA_LAUNCH_BLOCKING=1 for +immediate error reporting. +""" + +import pytest +import torch +import torch.nn.functional as F + +from fla.utils import device + +# Dimensions that trigger int32 overflow: B*T*H*K = 4096*576*8*128 = 2.4B > 2^31 +B, T, H, K, V = 4096, 576, 8, 128, 128 + + +def _has_enough_gpu_memory(min_gb=20): + """Check CUDA availability and allocate min_gb to confirm it's usable.""" + if not torch.cuda.is_available(): + return False + try: + x = torch.empty(int(min_gb * 1024**3 // 4), dtype=torch.float32, device='cuda') + del x + torch.cuda.empty_cache() + return True + except torch.cuda.OutOfMemoryError: + return False + + +requires_large_gpu = pytest.mark.skipif( + not _has_enough_gpu_memory(20), + reason='Requires CUDA with >= 20GB allocatable memory' +) + + +@requires_large_gpu +def test_gated_delta_rule_chunk(): + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = F.normalize(torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16), p=2, dim=-1) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + g = F.logsigmoid(torch.randn(B, T, H, device=device, dtype=torch.float32)) + beta = torch.rand(B, T, H, device=device, dtype=torch.float32).sigmoid() + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = chunk_gated_delta_rule(q, k, v, g=g, beta=beta, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_gated_delta_rule_fused_recurrent(): + from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = F.normalize(torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16), p=2, dim=-1) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + g = F.logsigmoid(torch.randn(B, T, H, device=device, dtype=torch.float32)) + beta = torch.rand(B, T, H, device=device, dtype=torch.float32).sigmoid() + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = fused_recurrent_gated_delta_rule(q, k, v, g=g, beta=beta, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_delta_rule_chunk(): + from fla.ops.delta_rule import chunk_delta_rule + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = F.normalize(torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16), p=2, dim=-1) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + beta = torch.rand(B, T, H, device=device, dtype=torch.float32).sigmoid() + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = chunk_delta_rule(q, k, v, beta=beta, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_gla_chunk(): + from fla.ops.gla import chunk_gla + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + g = F.logsigmoid(torch.randn(B, T, H, K, device=device, dtype=torch.float32)) + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = chunk_gla(q, k, v, g, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_hgrn_chunk(): + from fla.ops.hgrn import chunk_hgrn + D = 1024 + x = torch.randn(B, T, D, device=device, dtype=torch.float32) + g = torch.randn(B, T, D, device=device, dtype=torch.float32) + h0 = torch.randn(B, D, device=device, dtype=torch.float32) + o, ht = chunk_hgrn(x, g, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_rwkv6_chunk(): + from fla.ops.rwkv6 import chunk_rwkv6 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + w = torch.randn(B, T, H, K, device=device, dtype=torch.float32) + u = torch.randn(H, K, device=device, dtype=torch.float32) + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = chunk_rwkv6(q, k, v, w, u, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN" + + +@requires_large_gpu +def test_simple_gla_chunk(): + from fla.ops.simple_gla import chunk_simple_gla + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = F.normalize(torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16), p=2, dim=-1) + v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) + g = F.logsigmoid(torch.randn(B, T, H, device=device, dtype=torch.float32)) + h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) + o, ht = chunk_simple_gla(q, k, v, g=g, initial_state=h0, output_final_state=True) + torch.cuda.synchronize() + assert not o.isnan().any(), "Output contains NaN"