diff --git a/fla/layers/rwkv6.py b/fla/layers/rwkv6.py index 6d8f01660a..194aef58b1 100644 --- a/fla/layers/rwkv6.py +++ b/fla/layers/rwkv6.py @@ -160,10 +160,10 @@ def forward( ) elif mode == 'chunk': o, recurrent_state = chunk_rwkv6( - q=r, + r=r, k=k, v=v, - g=w, + w=w, u=u, scale=1., initial_state=recurrent_state, diff --git a/fla/ops/attn/parallel.py b/fla/ops/attn/parallel.py index 8e7b590c52..dce95717f1 100644 --- a/fla/ops/attn/parallel.py +++ b/fla/ops/attn/parallel.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +import warnings from typing import Optional import torch @@ -15,8 +16,8 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_G': lambda args: args['g_cumsum'] is not None + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -24,7 +25,7 @@ for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else []) for num_stages in [2, 3, 4, 5] ], - key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'], + key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'], ) @triton.jit def parallel_attn_fwd_kernel( @@ -35,8 +36,8 @@ def parallel_attn_fwd_kernel( g_cumsum, lse, scale, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -48,16 +49,16 @@ def parallel_attn_fwd_kernel( BS: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr ): i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + 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 else: i_n = i_b @@ -168,8 +169,8 @@ def parallel_attn_bwd_kernel_preprocess( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_G': lambda args: args['g_cumsum'] is not None + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -177,7 +178,7 @@ def parallel_attn_bwd_kernel_preprocess( for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else []) for num_stages in [2, 3, 4, 5] ], - key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'], + key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'], ) @triton.jit(do_not_specialize=['T']) def parallel_attn_bwd_kernel_dq( @@ -191,8 +192,8 @@ def parallel_attn_bwd_kernel_dq( dg_cumsum, g_cumsum, scale, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -212,8 +213,8 @@ def parallel_attn_bwd_kernel_dq( i_h = i_hq // G 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) + 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 else: i_n = i_b @@ -306,8 +307,8 @@ def parallel_attn_bwd_kernel_dq( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_G': lambda args: args['g_cumsum'] is not None + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -315,7 +316,7 @@ def parallel_attn_bwd_kernel_dq( for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else []) for num_stages in [2, 3, 4, 5] ], - key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'], + key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'], ) @triton.jit(do_not_specialize=['T']) def parallel_attn_bwd_kernel_dkv( @@ -329,8 +330,8 @@ def parallel_attn_bwd_kernel_dkv( dk, dv, dg_cumsum, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, B: tl.constexpr, @@ -351,8 +352,8 @@ def parallel_attn_bwd_kernel_dkv( i_h = i_hq // G 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) + 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 else: i_n = i_b @@ -464,7 +465,7 @@ def parallel_attn_fwd( g_cumsum: torch.Tensor, scale: float, chunk_size: int = 128, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] HQ = q.shape[2] @@ -485,8 +486,8 @@ def parallel_attn_fwd( NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) assert NK == 1, "The key dimension can not be larger than 256" o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) @@ -500,8 +501,8 @@ def parallel_attn_fwd( g_cumsum=g_cumsum, lse=lse, scale=scale, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, B=B, T=T, H=H, @@ -522,7 +523,7 @@ def parallel_attn_bwd_preprocess( do: torch.Tensor ): V = o.shape[-1] - delta = torch.empty_like(o[..., 0], dtype=torch.float32) + delta = torch.empty_like(o[..., 0], dtype=torch.float) parallel_attn_bwd_kernel_preprocess[(delta.numel(),)]( o=o, do=do, @@ -543,7 +544,7 @@ def parallel_attn_bwd( do: torch.Tensor, scale: float = None, chunk_size: int = 128, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] HQ = q.shape[2] @@ -554,8 +555,8 @@ def parallel_attn_bwd( BK = max(16, triton.next_power_of_2(K)) BV = max(16, triton.next_power_of_2(V)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NV = triton.cdiv(V, BV) delta = parallel_attn_bwd_preprocess(o, do) @@ -567,8 +568,8 @@ def parallel_attn_bwd( dg_cumsum, dg_cumsum_k = None, None if g_cumsum is not None: - dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) - dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) + dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) parallel_attn_bwd_kernel_dq[grid]( q=q, @@ -580,8 +581,8 @@ def parallel_attn_bwd( do=do, dq=dq, dg_cumsum=dg_cumsum, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, B=B, @@ -606,8 +607,8 @@ def parallel_attn_bwd( dk=dk, dv=dv, dg_cumsum=dg_cumsum_k, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, B=B, @@ -647,7 +648,7 @@ def forward(ctx, q, k, v, g, scale, cu_seqlens): g_cumsum=g_cumsum, scale=scale, chunk_size=chunk_size, - offsets=cu_seqlens, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, o, g_cumsum, lse) ctx.chunk_size = chunk_size @@ -670,12 +671,12 @@ def backward(ctx, do): do=do, scale=ctx.scale, chunk_size=ctx.chunk_size, - offsets=ctx.cu_seqlens, + cu_seqlens=ctx.cu_seqlens, ) if dg is not None: dg = chunk_global_cumsum(dg, cu_seqlens=ctx.cu_seqlens, reverse=True) - return dq.to(q), dk.to(k), dv.to(v), dg, None, None, None, None, None, None, None, None + return dq.to(q), dk.to(k), dv.to(v), dg, None, None def parallel_attn( @@ -711,14 +712,26 @@ def parallel_attn( o (torch.Tensor): Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`. """ - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None: - assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" if head_first: + warnings.warn( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead." + ) q, k, v = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v)) if g is not None: g = rearrange(g, 'b h t ... -> b t h ...') + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...]." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + o = ParallelAttentionFunction.apply(q, k, v, g, scale, cu_seqlens) if head_first: o = rearrange(o, 'b t h ... -> b h t ...') diff --git a/fla/ops/based/fused_chunk.py b/fla/ops/based/fused_chunk.py index b017af069b..1bba63a72d 100644 --- a/fla/ops/based/fused_chunk.py +++ b/fla/ops/based/fused_chunk.py @@ -27,7 +27,6 @@ def fused_chunk_based_fwd_kernel( BK: tl.constexpr, BV: tl.constexpr, ): - # indices i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) o_i = tl.arange(0, BT) diff --git a/fla/ops/common/chunk_delta_h.py b/fla/ops/common/chunk_delta_h.py index 9a7f1c508e..d0f5d2f149 100644 --- a/fla/ops/common/chunk_delta_h.py +++ b/fla/ops/common/chunk_delta_h.py @@ -16,10 +16,10 @@ @triton.heuristics({ 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None + 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -41,7 +41,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( h, h0, ht, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -53,13 +53,13 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( USE_G: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, ): 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -192,7 +192,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 'USE_G': lambda args: args['g'] is not None, 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -216,7 +216,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( dh, dv, dv2, - offsets, + cu_seqlens, chunk_offsets, scale, T, @@ -233,7 +233,7 @@ 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 // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -385,8 +385,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, 'USE_Q': lambda args: args['q'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -407,20 +407,20 @@ def proprocess_qkw( q_new, k_new, w_new, - offsets, + cu_seqlens, T, H: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, + USE_Q: tl.constexpr, IS_VARLEN: tl.constexpr, - USE_Q: tl.constexpr ): i_k, i_nh, i_t = 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 else: bos, eos = i_n * T, i_n * T + T @@ -472,19 +472,19 @@ def chunk_gated_delta_rule_fwd_h( g: Optional[torch.Tensor] = None, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, chunk_size: int = 64, # SY: remove this argument and force chunk size 64? save_new_value: bool = True, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, u.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) assert K <= 256, "current kernel does not support head dimension larger than 256." h = k.new_empty(B, NT, H, K, V) @@ -502,7 +502,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), N*H, triton.cdiv(T, BT)) q_new=None, k_new=k_new, w_new=w_new, - offsets=offsets, + cu_seqlens=cu_seqlens, T=T, H=H, K=K, @@ -520,7 +520,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) h=h, h0=initial_state, ht=final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, @@ -542,7 +542,7 @@ def chunk_gated_delta_rule_bwd_dhu( do: torch.Tensor, dv: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64, # SY: remove this argument and force chunk size 64? ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *q.shape, do.shape[-1] @@ -550,11 +550,11 @@ def chunk_gated_delta_rule_bwd_dhu( BT = 64 assert K <= 256, "current kernel does not support head dimension being larger than 256." - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - if offsets is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) dh = q.new_empty(B, NT, H, K, V) dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None @@ -573,7 +573,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), N*H, triton.cdiv(T, BT)) q_new=q_new, k_new=k_new, w_new=w_new, - offsets=offsets, + cu_seqlens=cu_seqlens, T=T, H=H, K=K, @@ -592,7 +592,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) dh=dh, dv=dv, dv2=dv2, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, scale=scale, T=T, diff --git a/fla/ops/common/chunk_h.py b/fla/ops/common/chunk_h.py index a72d5c23db..5c7dd0c447 100644 --- a/fla/ops/common/chunk_h.py +++ b/fla/ops/common/chunk_h.py @@ -17,7 +17,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -39,7 +39,7 @@ def chunk_fwd_kernel_h( gv, h0, ht, - offsets, + cu_seqlens, split_offsets, T, H: tl.constexpr, @@ -59,7 +59,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) @@ -132,7 +132,7 @@ def chunk_fwd_kernel_h( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -154,7 +154,7 @@ def chunk_bwd_kernel_dh( dh, dht, dh0, - offsets, + cu_seqlens, split_offsets, scale, T, @@ -178,7 +178,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) @@ -253,7 +253,7 @@ def chunk_fwd_h( gv: torch.Tensor, h0: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, chunk_size: int = 64, split_size: Optional[int] = None, states_in_fp32: bool = False @@ -263,11 +263,11 @@ def chunk_fwd_h( BS = BT if split_size is None else min(split_size, max(16, triton.next_power_of_2(T))) assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NS, split_offsets = B, triton.cdiv(T, BS), None else: - split_offsets = prepare_chunk_offsets(offsets, BS) - N, NS = len(offsets) - 1, split_offsets[-1].item() + split_offsets = prepare_chunk_offsets(cu_seqlens, BS) + N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None @@ -281,7 +281,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), gv=gv, h0=h0, ht=ht, - offsets=offsets, + cu_seqlens=cu_seqlens, split_offsets=split_offsets, T=T, H=H, @@ -307,7 +307,7 @@ def chunk_bwd_dh( h0: torch.Tensor, dht: torch.Tensor, scale: float, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, chunk_size: int = 64, split_size: Optional[int] = None, states_in_fp32: bool = False @@ -319,11 +319,11 @@ def chunk_bwd_dh( assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" # N: the actual number of sequences in the batch with either equal or variable lengths # NG: number of groups in GQA - if offsets is None: + if cu_seqlens is None: N, NS, split_offsets = B, triton.cdiv(T, BS), None else: - split_offsets = prepare_chunk_offsets(offsets, BS) - N, NS = len(offsets) - 1, split_offsets[-1].item() + split_offsets = prepare_chunk_offsets(cu_seqlens, BS) + N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() NG = HQ // H dh = k.new_empty(B, NS, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) @@ -339,7 +339,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), dh=dh, dht=dht, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, split_offsets=split_offsets, scale=scale, T=T, diff --git a/fla/ops/common/chunk_h_parallel.py b/fla/ops/common/chunk_h_parallel.py index 8bafbeec3b..cc0783b8a0 100644 --- a/fla/ops/common/chunk_h_parallel.py +++ b/fla/ops/common/chunk_h_parallel.py @@ -18,7 +18,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -40,8 +40,8 @@ def chunk_fwd_kernel_h_parallel( gv, h0, ht, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -68,8 +68,8 @@ def chunk_fwd_kernel_h_parallel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -132,7 +132,7 @@ def chunk_fwd_kernel_h_parallel( @triton.heuristics({ 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -152,7 +152,7 @@ def chunk_fwd_kernel_h_reduction( gv, kvt, ht, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -170,7 +170,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -217,7 +217,7 @@ def chunk_fwd_kernel_h_reduction( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -239,8 +239,8 @@ def chunk_bwd_kernel_dh_parallel( dh, dht, dh0, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, HQ: tl.constexpr, @@ -266,8 +266,8 @@ def chunk_bwd_kernel_dh_parallel( i_h = i_hq // NG if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -320,7 +320,7 @@ def chunk_bwd_kernel_dh_parallel( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -340,7 +340,7 @@ def chunk_bwd_kernel_dh_reduction( dh, doq0, dh0, - offsets, + cu_seqlens, chunk_offsets, T, HQ: tl.constexpr, @@ -361,7 +361,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -409,17 +409,18 @@ def chunk_fwd_h( h0: torch.Tensor, output_final_state: bool, states_in_fp32: bool = False, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - indices = prepare_chunk_indices(offsets, BT) - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) h = k.new_empty(B, NT, H, K, V, dtype=torch.float) ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None @@ -433,8 +434,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), gv=gv, h0=h0, ht=ht, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -453,7 +454,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), gv=gv, kvt=kvt, ht=ht, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, @@ -480,19 +481,20 @@ def chunk_bwd_dh( dht: torch.Tensor, scale: float, states_in_fp32: bool = False, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HQ = q.shape[2] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths # NG: number of groups in GQA - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - indices = prepare_chunk_indices(offsets, BT) - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) NG = HQ // H dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) @@ -508,8 +510,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), dh=dh, dht=dht, dh0=dh0, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, HQ=HQ, @@ -532,7 +534,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), dh=dh, doq0=doq0, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, HQ=HQ, diff --git a/fla/ops/common/chunk_h_split.py b/fla/ops/common/chunk_h_split.py index d270e27cc8..1a98f6be1c 100644 --- a/fla/ops/common/chunk_h_split.py +++ b/fla/ops/common/chunk_h_split.py @@ -13,7 +13,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -36,7 +36,7 @@ def chunk_fwd_kernel_h_split( hr, h0, ht, - offsets, + cu_seqlens, split_indices, T, S: tl.constexpr, @@ -61,7 +61,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NS = tl.cdiv(T, S) else: @@ -132,7 +132,7 @@ def chunk_fwd_kernel_h_split( @triton.heuristics({ 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -152,7 +152,7 @@ def chunk_fwd_kernel_h_reduction( hs, hr, ht, - offsets, + cu_seqlens, split_offsets, T, S: tl.constexpr, @@ -171,7 +171,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NS = tl.cdiv(T, S) boh = tl.load(split_offsets + i_n).to(tl.int32) @@ -218,7 +218,7 @@ def chunk_fwd_kernel_h_reduction( @triton.heuristics({ 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -241,7 +241,7 @@ def chunk_bwd_kernel_dh_split( dhs, dhr, dh0, - offsets, + cu_seqlens, split_indices, scale, T, @@ -269,7 +269,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NS = tl.cdiv(T, S) else: @@ -336,7 +336,7 @@ def chunk_bwd_kernel_dh_split( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -356,7 +356,7 @@ def chunk_bwd_kernel_dh_reduction( dhs, dhr, dh0, - offsets, + cu_seqlens, split_offsets, T, S: tl.constexpr, @@ -378,7 +378,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NS = tl.cdiv(T, S) boh = tl.load(split_offsets + i_n).to(tl.int32) @@ -427,7 +427,7 @@ def chunk_fwd_h( gv: torch.Tensor, h0: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, split_offsets: Optional[torch.LongTensor] = None, split_indices: Optional[torch.LongTensor] = None, chunk_size: int = 64, @@ -443,11 +443,11 @@ def chunk_fwd_h( # BT: chunk size S, BT = split_size, chunk_size assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" - if offsets is None: + if cu_seqlens is None: N = B NS = N * triton.cdiv(T, S) else: - N = len(offsets) - 1 + N = len(cu_seqlens) - 1 NS = split_offsets[-1] # unreduced kv states per split @@ -467,7 +467,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), hr=hr, h0=h0, ht=ht, - offsets=offsets, + cu_seqlens=cu_seqlens, split_indices=split_indices, T=T, S=S, @@ -487,7 +487,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), hs=hs, hr=hr, ht=ht, - offsets=offsets, + cu_seqlens=cu_seqlens, split_offsets=split_offsets, T=T, S=S, @@ -513,7 +513,7 @@ def chunk_bwd_dh( h0: torch.Tensor, dht: torch.Tensor, scale: float, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, split_offsets: Optional[torch.Tensor] = None, split_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, @@ -530,11 +530,11 @@ def chunk_bwd_dh( # BT: chunk size S, BT = max(chunk_size, min(split_size, triton.next_power_of_2(T))), chunk_size assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" - if offsets is None: + if cu_seqlens is None: N = B NS = N * triton.cdiv(T, S) else: - N = len(offsets) - 1 + N = len(cu_seqlens) - 1 NS = split_offsets[-1] # number of groups in GQA NG = HQ // H @@ -555,7 +555,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), dhs=dhs, dhr=dhr, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, split_indices=split_indices, scale=scale, T=T, @@ -579,7 +579,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), dhs=dhs, dhr=dhr, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, split_offsets=split_offsets, T=T, S=S, diff --git a/fla/ops/common/chunk_o.py b/fla/ops/common/chunk_o.py index dd96f498e8..46abf77e33 100644 --- a/fla/ops/common/chunk_o.py +++ b/fla/ops/common/chunk_o.py @@ -17,7 +17,7 @@ @triton.heuristics({ 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -37,8 +37,8 @@ def chunk_fwd_kernel_o( h, g, o, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -55,8 +55,8 @@ def chunk_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -112,7 +112,7 @@ def chunk_fwd_kernel_o( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_G': lambda args: args['g'] is not None, 'USE_DW': lambda args: args['dw'] is not None }) @@ -139,8 +139,8 @@ def chunk_bwd_kernel_dqkwg( w, dv, dw, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, B: tl.constexpr, T, @@ -160,8 +160,8 @@ def chunk_bwd_kernel_dqkwg( dg += i_k * B * H * T if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -280,7 +280,7 @@ def chunk_bwd_kernel_dqkwg( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_G': lambda args: args['g'] is not None, }) @triton.autotune( @@ -299,8 +299,8 @@ def chunk_bwd_kernel_dv( do, dv, dh, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -316,8 +316,8 @@ def chunk_bwd_kernel_dv( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -366,7 +366,7 @@ def chunk_bwd_kernel_dv( @triton.heuristics({ 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -383,8 +383,8 @@ def chunk_bwd_kernel_dv_local( g, do, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -399,8 +399,8 @@ def chunk_bwd_kernel_dv_local( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -445,13 +445,13 @@ def chunk_fwd_o( h: torch.Tensor, g: Optional[torch.Tensor] = None, # cumsum of log decay scale: Optional[float] = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *q.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) if scale is None: scale = k.shape[-1] ** -0.5 @@ -465,8 +465,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) h, g, o, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -484,12 +484,12 @@ def chunk_bwd_dv( do: torch.Tensor, dh: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *k.shape, do.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # H100 can have larger block size if check_shared_mem('hopper', k.device.index): CONST_TILING = 128 @@ -499,7 +499,7 @@ def chunk_bwd_dv( CONST_TILING = 32 BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NV = triton.cdiv(V, BV) dv = torch.empty_like(do) @@ -511,8 +511,8 @@ def chunk_bwd_dv( do, dv, dh, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -530,14 +530,13 @@ def chunk_bwd_dv_local( k: torch.Tensor, g: torch.Tensor, do: torch.Tensor, - dh: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *k.shape, do.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # H100 can have larger block size if check_shared_mem('hopper', k.device.index): CONST_TILING = 128 @@ -547,7 +546,7 @@ def chunk_bwd_dv_local( CONST_TILING = 32 BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dv = torch.empty_like(do) grid = (NT, B * H) @@ -557,8 +556,8 @@ def chunk_bwd_dv_local( g, do, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -581,15 +580,15 @@ def chunk_bwd_dqkwg( dh: torch.Tensor, dv: Optional[torch.Tensor] = None, w: Optional[torch.Tensor] = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64, scale: float = 1.0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) CONST_TILING = 64 if check_shared_mem() else 32 BK = min(triton.next_power_of_2(K), CONST_TILING) @@ -615,8 +614,8 @@ def chunk_bwd_dqkwg( dq=dq, dk=dk, dg=dg, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, B=B, T=T, diff --git a/fla/ops/common/chunk_scaled_dot_kkt.py b/fla/ops/common/chunk_scaled_dot_kkt.py index 22cc7a0109..69faf8c85c 100644 --- a/fla/ops/common/chunk_scaled_dot_kkt.py +++ b/fla/ops/common/chunk_scaled_dot_kkt.py @@ -11,7 +11,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -27,8 +27,8 @@ def chunk_scaled_dot_kkt_fwd_kernel( k, beta, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -39,8 +39,8 @@ def chunk_scaled_dot_kkt_fwd_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -89,15 +89,15 @@ def chunk_scaled_dot_kkt_fwd( """ B, T, H, K = k.shape BT = chunk_size - indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( k=k, beta=beta, A=A, - offsets=cu_seqlens, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, diff --git a/fla/ops/common/fused_recurrent.py b/fla/ops/common/fused_recurrent.py index 2c3ae800bc..9d75090b0c 100644 --- a/fla/ops/common/fused_recurrent.py +++ b/fla/ops/common/fused_recurrent.py @@ -15,7 +15,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -35,7 +35,7 @@ def fused_recurrent_fwd_kernel( o, h0, ht, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -52,11 +52,10 @@ def fused_recurrent_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr ): - # indices i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -120,7 +119,7 @@ def fused_recurrent_fwd_kernel( 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -144,7 +143,7 @@ def fused_recurrent_bwd_kernel( dv, dht, dh0, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -165,7 +164,7 @@ def fused_recurrent_bwd_kernel( i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -291,10 +290,10 @@ def fused_recurrent_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = min(K, 64), min(V, 64) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) @@ -313,7 +312,7 @@ def fused_recurrent_fwd( o, h0, ht, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -344,10 +343,10 @@ def fused_recurrent_bwd( scale: Optional[float] = None, initial_state: Optional[torch.Tensor] = None, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = min(K, 64), min(V, 64) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) @@ -373,7 +372,7 @@ def fused_recurrent_bwd( dv, dht, dh0, - offsets, + cu_seqlens, scale, B=B, T=T, @@ -392,11 +391,11 @@ def fused_recurrent_bwd( dv = dv.sum(0) dg, dgk, dgv = None, None, None if g is not None: - dg = chunk_global_cumsum((dq * q.float() - dk * k.float()).sum(-1), reverse=not reverse, cu_seqlens=offsets) + dg = chunk_global_cumsum((dq * q.float() - dk * k.float()).sum(-1), reverse=not reverse, cu_seqlens=cu_seqlens) if gk is not None: - dgk = chunk_global_cumsum(dq * q.float() - dk * k.float(), reverse=not reverse, cu_seqlens=offsets) + dgk = chunk_global_cumsum(dq * q.float() - dk * k.float(), reverse=not reverse, cu_seqlens=cu_seqlens) if gv is not None: - dgv = chunk_global_cumsum(do.float() * o.float() - dv * v.float(), reverse=not reverse, cu_seqlens=offsets) + dgv = chunk_global_cumsum(do.float() * o.float() - dv * v.float(), reverse=not reverse, cu_seqlens=cu_seqlens) return dq, dk, dv, dg, dgk, dgv, dh0 @@ -418,7 +417,7 @@ def forward( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ): o, ht = fused_recurrent_fwd( q=q, @@ -431,12 +430,12 @@ def forward( initial_state=initial_state, output_final_state=output_final_state, reverse=reverse, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, g, gk, gv, initial_state, o) ctx.scale = scale ctx.reverse = reverse - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o.to(q.dtype), ht @staticmethod @@ -466,7 +465,7 @@ def backward(ctx, do, dht): scale=ctx.scale, initial_state=initial_state, reverse=ctx.reverse, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, ) return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dg, dgk, dgv, None, dh0, None, None, None diff --git a/fla/ops/common/utils.py b/fla/ops/common/utils.py index c61cf9a36b..05396abc91 100644 --- a/fla/ops/common/utils.py +++ b/fla/ops/common/utils.py @@ -18,11 +18,11 @@ @triton.jit def prepare_position_ids_kernel( y, - offsets, + cu_seqlens, B: tl.constexpr ): i_n = tl.program_id(0) - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 o = tl.arange(0, B) @@ -32,13 +32,16 @@ def prepare_position_ids_kernel( @tensor_cache -def prepare_lens(offsets: torch.LongTensor) -> torch.LongTensor: - return offsets[1:] - offsets[:-1] +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return cu_seqlens[1:] - cu_seqlens[:-1] @tensor_cache -def prepare_position_ids(offsets: torch.LongTensor) -> torch.LongTensor: - return torch.cat([torch.arange(n, dtype=offsets.dtype, device=offsets.device) for n in prepare_lens(offsets).unbind()]) +def prepare_position_ids(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return torch.cat([ + torch.arange(n, dtype=cu_seqlens.dtype, device=cu_seqlens.device) + for n in prepare_lens(cu_seqlens).unbind() + ]) @tensor_cache @@ -47,23 +50,23 @@ def prepare_sequence_ids(position_ids: torch.LongTensor) -> torch.LongTensor: @tensor_cache -def prepare_token_indices(offsets: torch.LongTensor) -> torch.LongTensor: - position_ids = prepare_position_ids(offsets) - return torch.stack([prepare_sequence_ids(position_ids), position_ids], 1).to(offsets) +def prepare_token_indices(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + position_ids = prepare_position_ids(cu_seqlens) + return torch.stack([prepare_sequence_ids(position_ids), position_ids], 1).to(cu_seqlens) @tensor_cache def prepare_chunk_indices( - offsets: torch.LongTensor, + cu_seqlens: torch.LongTensor, chunk_size: int ) -> torch.LongTensor: - indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(offsets), chunk_size).tolist()]) - return torch.stack([prepare_sequence_ids(indices), indices], 1).to(offsets) + indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) + return torch.stack([prepare_sequence_ids(indices), indices], 1).to(cu_seqlens) @tensor_cache def prepare_chunk_offsets( - offsets: torch.LongTensor, + cu_seqlens: torch.LongTensor, chunk_size: int ) -> torch.LongTensor: - return torch.cat([offsets.new_tensor([0]), triton.cdiv(prepare_lens(offsets), chunk_size)]).cumsum(-1) + return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) diff --git a/fla/ops/delta_rule/chunk.py b/fla/ops/delta_rule/chunk.py index 7c31cdd82c..1c6a6d9f26 100644 --- a/fla/ops/delta_rule/chunk.py +++ b/fla/ops/delta_rule/chunk.py @@ -23,7 +23,7 @@ def chunk_delta_rule_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] @@ -33,7 +33,7 @@ def chunk_delta_rule_fwd( k=k, v=v, beta=beta, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -44,7 +44,7 @@ def chunk_delta_rule_fwd( g=None, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) o = chunk_fwd_o( @@ -54,7 +54,7 @@ def chunk_delta_rule_fwd( h=h, g=None, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return o, A, final_state @@ -70,7 +70,7 @@ def chunk_delta_rule_bwd( initial_state: torch.Tensor, do: torch.Tensor, dht: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] @@ -80,7 +80,7 @@ def chunk_delta_rule_bwd( v=v, beta=beta, A=A, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) h, v_new, _ = chunk_gated_delta_rule_fwd_h( @@ -90,7 +90,7 @@ def chunk_delta_rule_bwd( g=None, initial_state=initial_state, output_final_state=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dv = chunk_bwd_dv_local( @@ -98,9 +98,8 @@ def chunk_delta_rule_bwd( k=k, do=do, g=None, - dh=None, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( @@ -113,7 +112,7 @@ def chunk_delta_rule_bwd( do=do, dv=dv, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dq, dk, dw, _ = chunk_bwd_dqkwg( @@ -127,7 +126,7 @@ def chunk_delta_rule_bwd( dh=dh, g=None, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dk2, dv, db = bwd_prepare_wy_repr( @@ -137,7 +136,7 @@ def chunk_delta_rule_bwd( A=A, dw=dw, du=dv, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dk.add_(dk2) @@ -158,7 +157,7 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = True ): T = q.shape[1] @@ -179,13 +178,13 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) ctx.save_for_backward(q_orig, k_orig, v, beta, A, initial_state) ctx.chunk_size = chunk_size ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel return o.to(q.dtype), final_state @@ -213,7 +212,7 @@ def backward( initial_state=initial_state, do=do, dht=dht, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, chunk_size=ctx.chunk_size ) if use_qk_l2norm_in_kernel: diff --git a/fla/ops/delta_rule/fused_recurrent.py b/fla/ops/delta_rule/fused_recurrent.py index 72f17f72b6..b36b7a2fa5 100644 --- a/fla/ops/delta_rule/fused_recurrent.py +++ b/fla/ops/delta_rule/fused_recurrent.py @@ -16,7 +16,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.jit(do_not_specialize=['T']) def fused_recurrent_delta_rule_fwd_kernel( @@ -28,7 +28,7 @@ def fused_recurrent_delta_rule_fwd_kernel( o, h0, ht, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -45,7 +45,7 @@ def fused_recurrent_delta_rule_fwd_kernel( i_v, i_k, 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -103,7 +103,7 @@ def fused_recurrent_delta_rule_fwd_kernel( @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.jit(do_not_specialize=['T']) def fused_recurrent_delta_rule_bwd_kernel( @@ -119,7 +119,7 @@ def fused_recurrent_delta_rule_bwd_kernel( dk, dv, db, - offsets, + cu_seqlens, scale, B: tl.constexpr, T, @@ -137,7 +137,7 @@ def fused_recurrent_delta_rule_bwd_kernel( i_v, i_k, 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -261,10 +261,10 @@ def fused_recurrent_delta_rule_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" @@ -288,7 +288,7 @@ def fused_recurrent_delta_rule_fwd( o, initial_state, final_state, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -314,10 +314,10 @@ def fused_recurrent_delta_rule_bwd( do: torch.Tensor, scale: float, initial_state: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" @@ -353,7 +353,7 @@ def fused_recurrent_delta_rule_bwd( dk, dv, db, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -388,7 +388,7 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False ): q_orig = q @@ -406,12 +406,12 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q_orig, k_orig, u, beta, initial_state) ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel return o, final_state @@ -431,7 +431,7 @@ def backward(ctx, do, dht): do=do, scale=ctx.scale, initial_state=initial_state, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, ) if ctx.use_qk_l2norm_in_kernel: dq, dk = l2norm_bwd(q_orig, dq), l2norm_bwd(k_orig, dk) diff --git a/fla/ops/delta_rule/wy_fast.py b/fla/ops/delta_rule/wy_fast.py index a83439fe03..76af86f32d 100644 --- a/fla/ops/delta_rule/wy_fast.py +++ b/fla/ops/delta_rule/wy_fast.py @@ -16,7 +16,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -34,8 +34,8 @@ def fwd_recompute_w_u_kernel( w, u, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -48,8 +48,8 @@ def fwd_recompute_w_u_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -77,7 +77,7 @@ def fwd_recompute_w_u_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -98,8 +98,8 @@ def bwd_prepare_wy_repr_kernel( dk, dv, dbeta, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -112,8 +112,8 @@ def bwd_prepare_wy_repr_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -181,19 +181,19 @@ def fwd_prepare_wy_repr( k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A = chunk_scaled_dot_kkt_fwd( k=k, beta=beta, - cu_seqlens=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, output_dtype=torch.float32 ) A = solve_tril( A=A, - cu_seqlens=offsets, + cu_seqlens=cu_seqlens, output_dtype=k.dtype ) @@ -202,7 +202,7 @@ def fwd_prepare_wy_repr( v=v, beta=beta, A=A, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return w, u, A @@ -213,7 +213,7 @@ def fwd_recompute_w_u( v: torch.Tensor, beta: torch.Tensor, A: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] @@ -222,8 +222,8 @@ def fwd_recompute_w_u( BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) u = torch.empty_like(v) w = torch.empty_like(k) @@ -234,8 +234,8 @@ def fwd_recompute_w_u( w, u, A, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -254,7 +254,7 @@ def bwd_prepare_wy_repr( A: torch.Tensor, dw: torch.Tensor, du: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] @@ -263,8 +263,8 @@ def bwd_prepare_wy_repr( BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dk = torch.empty_like(k) dv = torch.empty_like(v) @@ -279,8 +279,8 @@ def bwd_prepare_wy_repr( dk, dv, dbeta, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, diff --git a/fla/ops/gated_delta_rule/chunk.py b/fla/ops/gated_delta_rule/chunk.py index 79b78fd051..81a18c7ecc 100644 --- a/fla/ops/gated_delta_rule/chunk.py +++ b/fla/ops/gated_delta_rule/chunk.py @@ -25,17 +25,17 @@ def chunk_gated_delta_rule_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): - g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=offsets) + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) # obtain WY representation. u is actually the new v. w, u, Aw, Au = fwd_prepare_wy_repr( k=k, v=v, beta=beta, g=g, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -46,7 +46,7 @@ def chunk_gated_delta_rule_fwd( g=g, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -58,7 +58,7 @@ def chunk_gated_delta_rule_fwd( h=h, g=g, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return g, o, Aw, Au, final_state @@ -76,7 +76,7 @@ def chunk_gated_delta_rule_bwd( initial_state: torch.Tensor, do: torch.Tensor, dht: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] @@ -87,7 +87,7 @@ def chunk_gated_delta_rule_bwd( beta=beta, Aw=Aw, Au=Au, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) h, v_new, _ = chunk_gated_delta_rule_fwd_h( @@ -97,7 +97,7 @@ def chunk_gated_delta_rule_bwd( g=g, initial_state=initial_state, output_final_state=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dv = chunk_bwd_dv_local( @@ -105,9 +105,8 @@ def chunk_gated_delta_rule_bwd( k=k, g=g, do=do, - dh=None, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( @@ -120,7 +119,7 @@ def chunk_gated_delta_rule_bwd( do=do, dv=dv, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dq, dk, dw, dg = chunk_bwd_dqkwg( @@ -134,7 +133,7 @@ def chunk_gated_delta_rule_bwd( do=do, dh=dh, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dk2, dv, db, dg2 = bwd_prepare_wy_repr( @@ -146,13 +145,13 @@ def chunk_gated_delta_rule_bwd( Au=Au, dw=dw, du=dv, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dk.add_(dk2) dg.add_(dg2) assert dg.dtype == torch.float32, "dg should be fp32" - dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=offsets) + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens) return dq, dk, dv, db, dg, dh0 @@ -171,7 +170,7 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False ): chunk_size = 64 @@ -191,10 +190,10 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, ) - ctx.save_for_backward(q_orig, k_orig, v, g, beta, Aw, Au, initial_state, offsets) + ctx.save_for_backward(q_orig, k_orig, v, g, beta, Aw, Au, initial_state, cu_seqlens) ctx.chunk_size = chunk_size ctx.scale = scale ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel @@ -208,7 +207,7 @@ def backward( do: torch.Tensor, dht: torch.Tensor ): - q, k, v, g, beta, Aw, Au, initial_state, offsets = ctx.saved_tensors + q, k, v, g, beta, Aw, Au, initial_state, cu_seqlens = ctx.saved_tensors if ctx.use_qk_l2norm_in_kernel: q, q_orig = l2norm_fwd(q), q k, k_orig = l2norm_fwd(k), k @@ -224,7 +223,7 @@ def backward( initial_state=initial_state, do=do, dht=dht, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=ctx.chunk_size ) if ctx.use_qk_l2norm_in_kernel: diff --git a/fla/ops/gated_delta_rule/fused_recurrent.py b/fla/ops/gated_delta_rule/fused_recurrent.py index ea1890a4d1..560a55cabe 100644 --- a/fla/ops/gated_delta_rule/fused_recurrent.py +++ b/fla/ops/gated_delta_rule/fused_recurrent.py @@ -16,7 +16,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.jit(do_not_specialize=['T']) def fused_recurrent_gated_delta_rule_fwd_kernel( @@ -28,7 +28,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( o, h0, ht, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -46,7 +46,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -121,10 +121,10 @@ def fused_recurrent_gated_delta_rule_fwd( initial_state: torch.Tensor, output_final_state: bool, use_qk_l2norm_in_kernel: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" @@ -147,7 +147,7 @@ def fused_recurrent_gated_delta_rule_fwd( o=o, h0=initial_state, ht=final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, scale=scale, T=T, B=B, @@ -179,7 +179,7 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False ): o, final_state = fused_recurrent_gated_delta_rule_fwd( @@ -192,7 +192,7 @@ def forward( initial_state=initial_state, output_final_state=output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - offsets=offsets + cu_seqlens=cu_seqlens ) return o, final_state diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py index 47c0892ba1..011623d12d 100644 --- a/fla/ops/gated_delta_rule/wy_fast.py +++ b/fla/ops/gated_delta_rule/wy_fast.py @@ -13,7 +13,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -30,8 +30,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( beta, Aw, Au, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -43,8 +43,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -87,7 +87,7 @@ def fwd_prepare_wy_repr_kernel_chunk32( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -104,8 +104,8 @@ def fwd_prepare_wy_repr_kernel_chunk64( beta, Aw, Au, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -117,8 +117,8 @@ def fwd_prepare_wy_repr_kernel_chunk64( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -203,7 +203,7 @@ def fwd_prepare_wy_repr_kernel_chunk64( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -222,8 +222,8 @@ def fwd_recompute_w_u_kernel( u, Aw, Au, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -236,8 +236,8 @@ def fwd_recompute_w_u_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -273,14 +273,14 @@ def fwd_prepare_wy_repr( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K = k.shape BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(BT, 32) BK = min(triton.next_power_of_2(K), 64) # bf16 should be good enough. @@ -294,8 +294,8 @@ def fwd_prepare_wy_repr( beta=beta, Aw=Aw, Au=Au, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -309,7 +309,7 @@ def fwd_prepare_wy_repr( beta=beta, Aw=Aw, Au=Au, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return w, u, Aw, Au @@ -321,14 +321,14 @@ def fwd_recompute_w_u( beta: torch.Tensor, Aw: torch.Tensor, Au: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = min(triton.next_power_of_2(K), 64) BV = min(triton.next_power_of_2(V), 64) @@ -342,8 +342,8 @@ def fwd_recompute_w_u( u=u, Aw=Aw, Au=Au, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -356,7 +356,7 @@ def fwd_recompute_w_u( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -380,8 +380,8 @@ def bwd_prepare_wy_repr_kernel( dv, dbeta, dg, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -394,8 +394,8 @@ def bwd_prepare_wy_repr_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -484,14 +484,14 @@ def bwd_prepare_wy_repr( Au: torch.Tensor, dw: torch.Tensor, du: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) CONST_TILING = 64 if check_shared_mem() else 32 BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) @@ -513,8 +513,8 @@ def bwd_prepare_wy_repr( dv=dv, dbeta=dbeta, dg=dg, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, diff --git a/fla/ops/generalized_delta_rule/dplr/chunk.py b/fla/ops/generalized_delta_rule/dplr/chunk.py index aad6afcf14..aaeaba3eeb 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk.py @@ -30,12 +30,12 @@ def chunk_dplr_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - gi, ge = chunk_rwkv6_fwd_cumsum(gk, BT, offsets=offsets) + gi, ge = chunk_rwkv6_fwd_cumsum(gk, BT, cu_seqlens=cu_seqlens) A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_fwd_intra_dplr_fn( q=q, @@ -45,7 +45,7 @@ def chunk_dplr_fwd( gi=gi, ge=ge, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, ) del ge @@ -57,7 +57,7 @@ def chunk_dplr_fwd( A_ab=A_ab, A_ak=A_ak, v=v, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del A_ab, A_ak @@ -70,7 +70,7 @@ def chunk_dplr_fwd( gk=gi, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del u, kg, bg, gi @@ -82,7 +82,7 @@ def chunk_dplr_fwd( A_qk=A_qk, A_qb=A_qb, h=h, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del v_new, h, A_qk, A_qb @@ -106,10 +106,9 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): chunk_size = 16 - o, final_state = chunk_dplr_fwd( q=q, k=k, @@ -120,11 +119,11 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) ctx.save_for_backward(q, k, v, a, b, gk, initial_state) - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.scale = scale ctx.chunk_size = chunk_size return o.to(q.dtype), final_state @@ -139,11 +138,11 @@ def backward( ): q, k, v, a, b, gk, initial_state = ctx.saved_tensors BT = ctx.chunk_size - offsets = ctx.offsets + cu_seqlens = ctx.cu_seqlens scale = ctx.scale # ******* start recomputing everything, otherwise i believe the gpu memory will be exhausted ******* - gi, ge = chunk_rwkv6_fwd_cumsum(gk, BT, offsets=offsets) + gi, ge = chunk_rwkv6_fwd_cumsum(gk, BT, cu_seqlens=cu_seqlens) A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_fwd_intra_dplr_fn( q=q, @@ -153,7 +152,7 @@ def backward( gi=gi, ge=ge, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, ) w, u, A_ab_inv = fwd_prepare_wy_repr( @@ -161,7 +160,7 @@ def backward( A_ab=A_ab, A_ak=A_ak, v=v, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del A_ab @@ -173,7 +172,7 @@ def backward( u=u, gk=gi, initial_state=initial_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del u @@ -187,7 +186,7 @@ def backward( do=do, A_qb=A_qb, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -200,7 +199,7 @@ def backward( dht=dht, do=do, dv=dv_new_intra, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -209,7 +208,7 @@ def backward( kg=kg, do=do, dh=dh, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del A_qk @@ -225,7 +224,7 @@ def backward( dv=dv_new, w=w, gk=gi, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, scale=scale, ) @@ -239,7 +238,7 @@ def backward( dw=dw, du=dv_new, dv0=dv, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) del A_ak @@ -262,7 +261,7 @@ def backward( dbg=dbg, chunk_size=BT, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, ) return dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), dgk.to(gk), None, dh0, None, None @@ -281,7 +280,6 @@ def chunk_dplr_delta_rule( output_final_state: bool = False, cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False, - input_precision: Optional[torch.dtype] = torch.bfloat16, ): r""" Args: @@ -312,9 +310,6 @@ def chunk_dplr_delta_rule( head_first (Optional[bool]): Whether the inputs are in the head-first format, which is not supported for variable-length inputs. Default: `False`. - input_precision (Optional[torch.dtype]): - The precision of the input tensors. Default: `torch.bfloat16`. - Note: The output tensors will be in the same precision as the input tensors. Use torch.float16 with caution. Returns: o (torch.Tensor): @@ -336,13 +331,6 @@ def chunk_dplr_delta_rule( "Please verify your input tensor format matches the expected shape [B, T, H, ...]." ) - # use pytorch fast path here, if q, k, v are already in input_precision, nothing to do - if input_precision == torch.float32: - warnings.warn( - """ChunkDeltaRuleFunction does not support float32. Please use bfloat16. - If you want to use float32, please solve the issue by yourself.""" - ) - q, k, v = q.to(input_precision), k.to(input_precision), v.to(input_precision) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( 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 dfabb67560..ccf8d0b33e 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py @@ -13,7 +13,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -46,8 +46,8 @@ def chunk_dplr_bwd_kernel_intra( dbg, dgk, dgk_offset, - offsets, - indices, + cu_seqlens, + chunk_indices, scale: tl.constexpr, T, H: tl.constexpr, @@ -63,8 +63,8 @@ def chunk_dplr_bwd_kernel_intra( i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC 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) + 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) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -289,7 +289,7 @@ def chunk_dplr_bwd_kernel_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -307,8 +307,8 @@ def chunk_dplr_bwd_dgk_kernel( dgk_offset, dgk_last, dgk_output, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -320,8 +320,8 @@ def chunk_dplr_bwd_dgk_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -366,8 +366,8 @@ def chunk_dplr_bwd_dqk_intra( dag: torch.Tensor, dbg: torch.Tensor, dgk_last: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, scale: float = 1.0, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64, ): B, T, H, K = q.shape @@ -375,8 +375,8 @@ def chunk_dplr_bwd_dqk_intra( BC = min(16, BT) BK = min(64, triton.next_power_of_2(K)) if check_shared_mem() else min(32, triton.next_power_of_2(K)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) @@ -409,8 +409,8 @@ def chunk_dplr_bwd_dqk_intra( dbg=dbg, da=da, db=db, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, H=H, @@ -422,16 +422,16 @@ def chunk_dplr_bwd_dqk_intra( GATHER_SUPPORTED=is_gather_supported ) - def grid2(meta): return (NT, triton.cdiv(K, meta['BK']), B * H) dgk_output = torch.empty_like(dgk) - chunk_dplr_bwd_dgk_kernel[grid2]( + def grid(meta): return (NT, triton.cdiv(K, meta['BK']), B * H) + chunk_dplr_bwd_dgk_kernel[grid]( dgk=dgk, dgk_offset=dgk_offset, dgk_last=dgk_last, dgk_output=dgk_output, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=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 b079e691e2..a6045f5085 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py @@ -13,7 +13,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -37,8 +37,8 @@ def chunk_dplr_fwd_A_kernel_intra_sub_inter( Aqb, Aab, Aak, - offsets, - indices, + cu_seqlens, + chunk_indices, scale: tl.constexpr, T, H: tl.constexpr, @@ -53,8 +53,8 @@ def chunk_dplr_fwd_A_kernel_intra_sub_inter( i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -113,7 +113,7 @@ def chunk_dplr_fwd_A_kernel_intra_sub_inter( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -140,8 +140,8 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( Aqb, Aab, Aak, - offsets, - indices, + cu_seqlens, + chunk_indices, scale: tl.constexpr, T, H: tl.constexpr, @@ -157,8 +157,8 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( i_b, i_h = i_bh // H, i_bh % H i_j = i_i 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -247,13 +247,13 @@ def chunk_fwd_intra_dplr_fn( ge: torch.Tensor, scale: float, chunk_size: int, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K = k.shape BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(16, BT) NC = triton.cdiv(BT, BC) @@ -275,8 +275,8 @@ def chunk_fwd_intra_dplr_fn( Aqb=Aqb, Aab=Aab, Aak=Aak, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, H=H, @@ -306,8 +306,8 @@ def chunk_fwd_intra_dplr_fn( kg=kg, ag=ag, bg=bg, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, H=H, 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 071b155dbf..df068e4851 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py @@ -15,7 +15,7 @@ @triton.heuristics({ 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -38,7 +38,7 @@ def chunk_dplr_bwd_kernel_dhu( dh, dv, dv2, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -55,7 +55,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -114,7 +114,7 @@ def chunk_dplr_bwd_dhu( dht: Optional[torch.Tensor], do: torch.Tensor, dv: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *qg.shape, do.shape[-1] @@ -132,12 +132,12 @@ def chunk_dplr_bwd_dhu( BV = 16 BC = 16 - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BC = min(BT, BC) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) @@ -159,7 +159,7 @@ def chunk_dplr_bwd_dhu( dh=dh, dv=dv, dv2=dv2, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, 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 1e4adaeba6..3f278ca672 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py @@ -15,7 +15,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -38,7 +38,7 @@ def chunk_dplr_fwd_kernel_h( h, h0, ht, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -56,7 +56,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -114,17 +114,18 @@ def chunk_dplr_fwd_h( gk: torch.Tensor, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *kg.shape, u.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) + + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - indices = prepare_chunk_indices(offsets, BT) - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BK = triton.next_power_of_2(K) assert BK <= 256, "current kernel does not support head dimension larger than 256." # H100 can have larger block size @@ -159,7 +160,7 @@ def chunk_dplr_fwd_h( gk=gk, h0=initial_state, ht=final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, 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 1cc82d0d84..29bf770b4e 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py @@ -15,7 +15,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -35,8 +35,8 @@ def chunk_dplr_bwd_kernel_dAu( dA_qk, dA_qb, dv_new, - offsets, - indices, + cu_seqlens, + chunk_indices, scale: tl.constexpr, T, H: tl.constexpr, @@ -48,8 +48,8 @@ def chunk_dplr_bwd_kernel_dAu( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -87,7 +87,7 @@ def chunk_dplr_bwd_kernel_dAu( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -115,8 +115,8 @@ def chunk_dplr_bwd_o_kernel( dgk_last, k, b, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -131,8 +131,8 @@ def chunk_dplr_bwd_o_kernel( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -214,7 +214,7 @@ def chunk_dplr_bwd_o_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -234,8 +234,8 @@ def chunk_dplr_bwd_kernel_dv( do, dv, dh, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -249,8 +249,8 @@ def chunk_dplr_bwd_kernel_dv( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -292,14 +292,14 @@ def chunk_dplr_bwd_dv( kg: torch.Tensor, do: torch.Tensor, dh: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *kg.shape, do.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dv = torch.empty_like(do) @@ -310,8 +310,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) do=do, dv=dv, dh=dh, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -332,7 +332,7 @@ def chunk_dplr_bwd_o( dh: torch.Tensor, dv: torch.Tensor, w: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64, scale: float = 1.0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -340,8 +340,8 @@ def chunk_dplr_bwd_o( B, T, H, K, V = *w.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = min(triton.next_power_of_2(K), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) BV = min(triton.next_power_of_2(V), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) @@ -370,8 +370,8 @@ def chunk_dplr_bwd_o( dv=dv, dw=dw, gk=gk, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -389,13 +389,13 @@ def chunk_dplr_bwd_dAu( do: torch.Tensor, A_qb: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, V = v.shape BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) if check_shared_mem('ampere'): # A100 BV = min(triton.next_power_of_2(V), 128) @@ -416,8 +416,8 @@ def chunk_dplr_bwd_dAu( dA_qk=dA_qk, dA_qb=dA_qb, dv_new=dv_new, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, H=H, 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 c46379cdba..f1548a6b9a 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py @@ -14,7 +14,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -36,8 +36,8 @@ def chunk_dplr_fwd_kernel_o( A_qb, h, o, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -52,8 +52,8 @@ def chunk_dplr_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -93,14 +93,14 @@ def chunk_dplr_fwd_o( A_qk: torch.Tensor, A_qb: torch.Tensor, h: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *qg.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = torch.empty_like(v) def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) @@ -112,8 +112,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) A_qb=A_qb, h=h, o=o, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, diff --git a/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py index eaace6e087..ca3e21d668 100644 --- a/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py +++ b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py @@ -16,7 +16,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -39,7 +39,7 @@ def fused_recurrent_dplr_delta_rule_fwd_kernel( o, h0, ht, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -57,7 +57,7 @@ def fused_recurrent_dplr_delta_rule_fwd_kernel( i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).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 else: bos, eos = i_n * T, i_n * T + T @@ -118,10 +118,10 @@ def fused_recurrent_dplr_delta_rule_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK = triton.next_power_of_2(K) h0 = initial_state @@ -142,7 +142,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) o, h0, ht, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -172,7 +172,7 @@ def forward( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): o, ht = fused_recurrent_dplr_delta_rule_fwd( q=q, @@ -185,7 +185,7 @@ def forward( initial_state=initial_state, output_final_state=output_final_state, reverse=reverse, - offsets=offsets, + cu_seqlens=cu_seqlens, ) return o, ht @@ -213,7 +213,6 @@ def fused_recurrent_dplr_delta_rule( reverse: bool = False, cu_seqlens: Optional[torch.Tensor] = None, head_first: bool = False, - input_precision: Optional[torch.dtype] = torch.bfloat16, ) -> Tuple[torch.Tensor, torch.Tensor]: r""" This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. @@ -248,9 +247,6 @@ def fused_recurrent_dplr_delta_rule( head_first (Optional[bool]): Whether the inputs are in the head-first format, which is not supported for variable-length inputs. Default: `False`. - input_precision (Optional[torch.dtype]): - The precision of the input tensors. Default: `torch.bfloat16`. - Note: The output tensors will be in the same precision as the input tensors. Use torch.float16 with caution. """ if head_first: warnings.warn( @@ -265,8 +261,7 @@ def fused_recurrent_dplr_delta_rule( "when head_first=False was specified. " "Please verify your input tensor format matches the expected shape [B, T, H, ...]." ) - # use pytorch fast path here, if q, k, v are already in input_precision, nothing to do - q, k, v = q.to(input_precision), k.to(input_precision), v.to(input_precision) + if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( 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 189d18ac05..f9ff4eb373 100644 --- a/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py @@ -15,12 +15,12 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ triton.Config(triton_config, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16, 32] + for num_warps in [2, 4, 8, 16] for num_stages in [2, 3, 4] ], key=['BT', 'BK', 'BV'], @@ -39,8 +39,8 @@ def bwd_prepare_wy_repr_kernel( dag, dAak, dAab, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -53,8 +53,8 @@ def bwd_prepare_wy_repr_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -121,15 +121,15 @@ def chunk_dplr_bwd_wy( dw: torch.Tensor, du: torch.Tensor, dv0: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A_ab_inv, A_ak, v, ag, dw, du = map(lambda x: x.contiguous(), [A_ab_inv, A_ak, v, ag, dw, du]) B, T, H, K, V = *dw.shape, du.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = min(triton.next_power_of_2(K), 64) BV = min(triton.next_power_of_2(V), 64) if check_shared_mem() else min(triton.next_power_of_2(V), 32) @@ -150,8 +150,8 @@ def chunk_dplr_bwd_wy( dag=dag, dAak=dA_ak, dAab=dA_ab, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, 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 cfbb2bc1d6..9699a4d2b0 100644 --- a/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py @@ -13,7 +13,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -27,8 +27,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( A_ab, A_ab_inv, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, BT: tl.constexpr, @@ -38,8 +38,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -57,7 +57,7 @@ def fwd_prepare_wy_repr_kernel_chunk32( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -72,8 +72,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( def fwd_prepare_wy_repr_kernel_chunk64( A_ab, A_ab_inv, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, BT: tl.constexpr, @@ -84,8 +84,8 @@ def fwd_prepare_wy_repr_kernel_chunk64( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -136,27 +136,27 @@ def fwd_prepare_wy_repr_kernel_chunk64( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16, 32] + for num_warps in [2, 4, 8, 16] for num_stages in [2, 3, 4] ], - key=['BT', 'BK', 'BV'], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], use_cuda_graph=use_cuda_graph, ) @triton.jit(do_not_specialize=['T']) def fwd_wu_kernel( - u, w, + u, ag, v, A_ab_inv, A_ak, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -169,17 +169,18 @@ def fwd_wu_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * 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)) p_A_ak = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_Aab_inv = tl.load(p_A_ab_inv, boundary_check=(0, 1)) b_Aak = tl.load(p_A_ak, boundary_check=(0, 1)) - o_s = tl.arange(0, BT) b_Aab_inv = tl.where(o_s[:, None] >= o_s[None, :], b_Aab_inv, 0) b_Aak = tl.where(o_s[:, None] > o_s[None, :], b_Aak, 0) # let's use tf32 here @@ -208,19 +209,19 @@ def fwd_wu( v: torch.Tensor, A_ak: torch.Tensor, A_ab_inv: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *ag.shape, v.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = min(triton.next_power_of_2(K), 64) BV = min(triton.next_power_of_2(V), 64) - u = torch.empty_like(v) w = torch.empty_like(ag) + u = torch.empty_like(v) fwd_wu_kernel[(NT, B * H)]( ag=ag, v=v, @@ -228,8 +229,8 @@ def fwd_wu( A_ab_inv=A_ab_inv, w=w, u=u, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -246,22 +247,22 @@ def fwd_prepare_wy_repr( v: torch.Tensor, A_ak: torch.Tensor, A_ab: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, _ = ag.shape BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(BT, 32) fwd_fn = fwd_prepare_wy_repr_kernel_chunk64 if BT == 64 else fwd_prepare_wy_repr_kernel_chunk32 A_ab_inv = torch.empty_like(A_ab) fwd_fn[(NT, B * H)]( A_ab=A_ab, A_ab_inv=A_ab_inv, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, BT=BT, @@ -272,7 +273,7 @@ def fwd_prepare_wy_repr( v=v, A_ak=A_ak, A_ab_inv=A_ab_inv, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return w, u, A_ab_inv diff --git a/fla/ops/generalized_delta_rule/iplr/chunk.py b/fla/ops/generalized_delta_rule/iplr/chunk.py index ad2c865355..ff0c6fd48a 100644 --- a/fla/ops/generalized_delta_rule/iplr/chunk.py +++ b/fla/ops/generalized_delta_rule/iplr/chunk.py @@ -19,7 +19,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -40,7 +40,7 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_h( h, h0, ht, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -58,7 +58,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -102,7 +102,7 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_h( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -124,8 +124,8 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_o( b, h, o, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -141,8 +141,8 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -206,7 +206,7 @@ def chunk_generalized_iplr_delta_rule_fwd_o( b: torch.Tensor, h: torch.Tensor, scale: Optional[float] = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *q.shape, v.shape[-1] @@ -214,8 +214,8 @@ def chunk_generalized_iplr_delta_rule_fwd_o( scale = k.shape[-1] ** -0.5 BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = torch.empty_like(v) @@ -232,8 +232,8 @@ def grid(meta): return ( b=b, h=h, o=o, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, H=H, @@ -252,18 +252,18 @@ def chunk_generalized_iplr_delta_rule_fwd_h( b: torch.Tensor, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, u.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BK = triton.next_power_of_2(K) assert BK <= 256, "current kernel does not support head dimension larger than 256." @@ -301,7 +301,7 @@ def chunk_generalized_iplr_delta_rule_fwd_h( h=h, h0=initial_state, ht=final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, @@ -325,7 +325,7 @@ def chunk_generalized_iplr_delta_rule_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] @@ -335,7 +335,7 @@ def chunk_generalized_iplr_delta_rule_fwd( b=b, k=k, v=v, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -347,7 +347,7 @@ def chunk_generalized_iplr_delta_rule_fwd( u=u, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) o = chunk_generalized_iplr_delta_rule_fwd_o( @@ -358,7 +358,7 @@ def chunk_generalized_iplr_delta_rule_fwd( b=b, h=h, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return o, final_state @@ -379,7 +379,7 @@ def forward( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): chunk_size = 64 @@ -392,7 +392,7 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return o.to(q.dtype), final_state diff --git a/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py index 5dc295675a..b9730adfa8 100644 --- a/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py +++ b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py @@ -15,7 +15,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -37,7 +37,7 @@ def fused_recurrent_fwd_kernel( ha, # tmp variable [B, H, L, V] for storing intermediate results of (h * a[None, :]).sum(0) h0, # initial hidden state [B, H, K, V] ht, # final hidden state [B, H, K, V] - offsets, # varlen offsets + cu_seqlens, # varlen cu_seqlens scale, # K ** -0.5 H, # n_heads T, # seq_len @@ -49,12 +49,11 @@ def fused_recurrent_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, # whether to store final state IS_VARLEN: tl.constexpr, ): - # indices 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).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 else: bos, eos = i_n * T, i_n * T + T @@ -107,7 +106,7 @@ def fused_recurrent_fwd_kernel( 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_DHT': lambda args: args['dht'] is not None, 'USE_DH0': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -138,7 +137,7 @@ def fused_recurrent_bwd_kernel( dha, # gradient of ha [NK, B, H, L, V] h0, # initial state [B, H, K, V] scale, # K ** -0.5 - offsets, # offsets + cu_seqlens, # cu_seqlens B, # batch_size H, # n_heads T, # seq_len @@ -158,7 +157,7 @@ def fused_recurrent_bwd_kernel( dq += i_v * B * H * K * T da += i_v * B * H * K * T if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).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 else: bos, eos = i_n * T, i_n * T + T @@ -281,9 +280,20 @@ class FusedRecurrentIPLRDeltaRuleFunction(torch.autograd.Function): @staticmethod @input_guard - def forward(ctx, q, k, v, a, b, scale=None, initial_state=None, output_final_state=False, offsets=None): + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None + ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK = triton.next_power_of_2(K) if output_final_state: @@ -309,7 +319,7 @@ def grid(meta): return ( h0=initial_state, ht=final_state, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, H=H, T=T, K=K, @@ -318,7 +328,7 @@ def grid(meta): return ( ) ctx.save_for_backward(q, k, v, a, b, ha, initial_state) ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o, final_state @staticmethod @@ -326,11 +336,10 @@ def grid(meta): return ( def backward(ctx, do, dht): q, k, v, a, b, ha, initial_state = ctx.saved_tensors B, T, H, K, V = *q.shape, v.shape[-1] - - N = B if ctx.offsets is None else len(ctx.offsets) - 1 - scale = ctx.scale + N = B if ctx.cu_seqlens is None else len(ctx.cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 64) NV = triton.cdiv(V, BV) + scale = ctx.scale dq = q.new_empty(NV, *q.shape) dk = k.new_empty(NV, *k.shape) @@ -363,7 +372,7 @@ def backward(ctx, do, dht): dha=dha, h0=initial_state, scale=scale, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, B=B, H=H, T=T, @@ -433,7 +442,7 @@ def fused_recurrent_iplr_delta_rule( if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `offsets`." + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." f"Please flatten variable-length inputs before processing." ) if head_first: diff --git a/fla/ops/generalized_delta_rule/iplr/wy_fast.py b/fla/ops/generalized_delta_rule/iplr/wy_fast.py index 7df6a7eca0..3f6b67899b 100644 --- a/fla/ops/generalized_delta_rule/iplr/wy_fast.py +++ b/fla/ops/generalized_delta_rule/iplr/wy_fast.py @@ -15,7 +15,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -29,8 +29,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( a, b, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -42,8 +42,8 @@ def fwd_prepare_wy_repr_kernel_chunk32( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -69,7 +69,7 @@ def fwd_prepare_wy_repr_kernel_chunk32( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -83,8 +83,8 @@ def fwd_prepare_wy_repr_kernel_chunk64( a, b, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -96,8 +96,8 @@ def fwd_prepare_wy_repr_kernel_chunk64( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -149,7 +149,7 @@ def fwd_prepare_wy_repr_kernel_chunk64( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -166,8 +166,8 @@ def fwd_wu_kernel( k, v, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -180,8 +180,8 @@ def fwd_wu_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -218,14 +218,14 @@ def fwd_wu( v: torch.Tensor, k: torch.Tensor, A: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *a.shape, v.shape[-1] BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) CONST_TILING = 64 if check_shared_mem() else 32 BK = min(triton.next_power_of_2(K), CONST_TILING) BV = min(triton.next_power_of_2(V), CONST_TILING) @@ -239,8 +239,8 @@ def fwd_wu( u=u, A=A, k=k, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -257,14 +257,14 @@ def fwd_prepare_wy_repr( b: torch.Tensor, v: torch.Tensor, k: torch.Tensor, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K = a.shape BT = min(chunk_size, max(triton.next_power_of_2(T), 16)) - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(BT, 32) BK = min(triton.next_power_of_2(K), 64) @@ -275,8 +275,8 @@ def fwd_prepare_wy_repr( a=a, b=b, A=A, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, K=K, @@ -289,7 +289,7 @@ def fwd_prepare_wy_repr( v=v, k=k, A=A, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return w, u, A diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index d0451bba18..16084e6edb 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -20,7 +20,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -37,8 +37,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( k, g, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -53,8 +53,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -93,7 +93,7 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -110,8 +110,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( k, g, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -125,8 +125,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( i_b, i_h = i_bh // H, i_bh % H i_j = i_i 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -158,7 +158,7 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -175,8 +175,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( k, g, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, B: tl.constexpr, @@ -193,8 +193,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( i_t, i_i = i_tc // NC, i_tc % NC i_j = i_i 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) + 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) all = T T = eos - bos else: @@ -229,7 +229,7 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -244,8 +244,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( A, A2, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -257,8 +257,8 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H 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) + 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) all = T T = eos - bos else: @@ -277,7 +277,7 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -296,8 +296,8 @@ def chunk_gla_fwd_kernel_o( h, o, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -312,8 +312,8 @@ def chunk_gla_fwd_kernel_o( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -355,7 +355,7 @@ def chunk_gla_fwd_kernel_o( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -372,8 +372,8 @@ def chunk_gla_bwd_kernel_intra( dA, dq, dk, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -387,8 +387,8 @@ def chunk_gla_bwd_kernel_intra( i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC 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) + 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) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -491,7 +491,7 @@ def chunk_gla_bwd_kernel_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -507,8 +507,8 @@ def chunk_gla_bwd_kernel_dA( v, do, dA, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -520,8 +520,8 @@ def chunk_gla_bwd_kernel_dA( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -540,7 +540,7 @@ def chunk_gla_bwd_kernel_dA( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -559,8 +559,8 @@ def chunk_gla_bwd_kernel_dv( do, dh, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -574,8 +574,8 @@ def chunk_gla_bwd_kernel_dv( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -614,7 +614,7 @@ def chunk_gla_bwd_kernel_dv( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -639,8 +639,8 @@ def chunk_gla_bwd_kernel_inter( dq2, dk2, dg, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -655,8 +655,8 @@ def chunk_gla_bwd_kernel_inter( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -723,14 +723,14 @@ def chunk_gla_fwd_intra_gk( k: torch.Tensor, g: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K = k.shape BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(16, BT) NC = triton.cdiv(BT, BC) @@ -741,8 +741,8 @@ def chunk_gla_fwd_intra_gk( k, g, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -761,8 +761,8 @@ def chunk_gla_fwd_intra_gk( k, g, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -783,8 +783,8 @@ def chunk_gla_fwd_intra_gk( k, g, A_intra, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, B=B, @@ -800,8 +800,8 @@ def chunk_gla_fwd_intra_gk( chunk_gla_fwd_A_kernel_intra_sub_intra_merge[grid]( A_intra, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, B=B, H=H, @@ -819,14 +819,14 @@ def chunk_gla_fwd_o_gk( A: torch.Tensor, h: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K, V = *q.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = torch.empty_like(v) def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) @@ -837,8 +837,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) h, o, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -853,14 +853,14 @@ def chunk_gla_bwd_dA( v: torch.Tensor, do: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, V = v.shape BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BV = min(64, triton.next_power_of_2(V)) dA = v.new_empty(B, T, H, BT, dtype=torch.float) @@ -869,8 +869,8 @@ def chunk_gla_bwd_dA( v, do, dA, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -887,14 +887,14 @@ def chunk_gla_bwd_dv( A: torch.Tensor, do: torch.Tensor, dh: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K, V = *k.shape, do.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dv = torch.empty_like(do) def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) @@ -905,8 +905,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) do, dh, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, K=K, @@ -921,7 +921,7 @@ def chunk_gla_bwd_dqk_intra( k: torch.Tensor, g: torch.Tensor, dA: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K = q.shape @@ -929,8 +929,8 @@ def chunk_gla_bwd_dqk_intra( BC = min(16, BT) BK = min(64, triton.next_power_of_2(K)) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) @@ -944,8 +944,8 @@ def chunk_gla_bwd_dqk_intra( dA, dq, dk, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, K=K, @@ -968,14 +968,14 @@ def chunk_gla_bwd_dqkg( dq: torch.Tensor, dk: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dg = torch.empty_like(g) dq2 = torch.empty_like(dq) @@ -994,8 +994,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) dq2, dk2, dg, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -1015,13 +1015,13 @@ def chunk_gla_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: T = q.shape[1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) if g_cumsum is None: - g_cumsum = chunk_local_cumsum(g, chunk_size=BT, cu_seqlens=offsets) + g_cumsum = chunk_local_cumsum(g, BT, cu_seqlens=cu_seqlens) h, ht = chunk_fwd_h( k=k, @@ -1032,7 +1032,7 @@ def chunk_gla_fwd( h0=initial_state, output_final_state=output_final_state, states_in_fp32=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -1043,7 +1043,7 @@ def chunk_gla_fwd( k=k, g=g_cumsum, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) o = chunk_gla_fwd_o_gk( @@ -1053,7 +1053,7 @@ def chunk_gla_fwd( A=A, h=h, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return g_cumsum, A, h, ht, o @@ -1071,13 +1071,13 @@ def chunk_gla_bwd( A: torch.Tensor, do: torch.Tensor, dht: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): T = q.shape[1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) if g_cumsum is None: - g_cumsum = chunk_local_cumsum(g, chunk_size=BT, cu_seqlens=offsets) + g_cumsum = chunk_local_cumsum(g, BT, cu_seqlens=cu_seqlens) if h is None: h, _ = chunk_fwd_h( @@ -1088,7 +1088,7 @@ def chunk_gla_bwd( gv=None, h0=initial_state, output_final_state=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, states_in_fp32=True ) @@ -1103,7 +1103,7 @@ def chunk_gla_bwd( h0=initial_state, dht=dht, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, states_in_fp32=True ) @@ -1114,7 +1114,7 @@ def chunk_gla_bwd( A=A, do=do, dh=dh, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) @@ -1123,7 +1123,7 @@ def chunk_gla_bwd( v=v, do=do, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dq, dk = chunk_gla_bwd_dqk_intra( @@ -1131,7 +1131,7 @@ def chunk_gla_bwd( k=k, g=g_cumsum, dA=dA, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dq, dk, dg = chunk_gla_bwd_dqkg( @@ -1145,7 +1145,7 @@ def chunk_gla_bwd( dq=dq, dk=dk, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return dq, dk, dv, dg, dh0 @@ -1164,7 +1164,7 @@ def forward( scale, initial_state, output_final_state, - offsets, + cu_seqlens, ): T = q.shape[1] chunk_size = min(64, max(16, triton.next_power_of_2(T))) @@ -1178,7 +1178,7 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) # recompute g_cumsum in bwd pass @@ -1189,14 +1189,14 @@ def forward( ctx.save_for_backward(q, k, v, g, g_cumsum, initial_state, A) ctx.chunk_size = chunk_size ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o, ht @staticmethod @input_guard def backward(ctx, do, dht): q, k, v, g, g_cumsum, initial_state, A = ctx.saved_tensors - chunk_size, scale, offsets = ctx.chunk_size, ctx.scale, ctx.offsets + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens dq, dk, dv, dg, dh0 = chunk_gla_bwd( q=q, k=k, @@ -1209,7 +1209,7 @@ def backward(ctx, do, dht): initial_state=initial_state, do=do, dht=dht, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq.to(q), dk.to(k), dv.to(v), dg, None, dh0, None, None diff --git a/fla/ops/gla/fused_chunk.py b/fla/ops/gla/fused_chunk.py index 584fab62a0..dc27f0fddf 100644 --- a/fla/ops/gla/fused_chunk.py +++ b/fla/ops/gla/fused_chunk.py @@ -131,7 +131,6 @@ def fused_chunk_gla_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, CHECK: tl.constexpr ): - # indices i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) b_h = tl.zeros([BK, BV], dtype=tl.float32) diff --git a/fla/ops/gsa/chunk.py b/fla/ops/gsa/chunk.py index c40ec33d4a..bdcea8907d 100644 --- a/fla/ops/gsa/chunk.py +++ b/fla/ops/gsa/chunk.py @@ -18,7 +18,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -38,8 +38,8 @@ def chunk_gsa_fwd_k_kernel_inter( g, o, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, HQ: tl.constexpr, @@ -57,8 +57,8 @@ def chunk_gsa_fwd_k_kernel_inter( i_h = i_hq // NG if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -102,7 +102,7 @@ def chunk_gsa_fwd_k_kernel_inter( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.jit(do_not_specialize=['T']) def chunk_gsa_fwd_k_kernel_intra( @@ -110,8 +110,8 @@ def chunk_gsa_fwd_k_kernel_intra( g, o, A, - offsets, - indices, + cu_seqlens, + chunk_indices, T, HQ: tl.constexpr, H: tl.constexpr, @@ -128,8 +128,8 @@ def chunk_gsa_fwd_k_kernel_intra( i_h = i_hq // NG i_t, i_i = i_c // NC, i_c % NC 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -182,7 +182,7 @@ def chunk_gsa_fwd_k_kernel_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -197,8 +197,8 @@ def chunk_gsa_bwd_k_kernel_dA( g, do, dA, - indices, - offsets, + chunk_indices, + cu_seqlens, scale, T, B: tl.constexpr, @@ -217,8 +217,8 @@ def chunk_gsa_bwd_k_kernel_dA( i_h = i_hq // NG 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(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) + 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) all = T T = eos - bos else: @@ -281,7 +281,7 @@ def chunk_gsa_bwd_k_kernel_dA( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -307,8 +307,8 @@ def chunk_gsa_bwd_k_kernel_dqkvg( dg, dgv, dA, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, B: tl.constexpr, @@ -327,8 +327,8 @@ def chunk_gsa_bwd_k_kernel_dqkvg( i_h = i_hq // NG if IS_VARLEN: i_tg = i_t - 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) + 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) all = T T = eos - bos NT = tl.cdiv(T, BT) @@ -415,7 +415,7 @@ def chunk_gsa_bwd_k_kernel_dqkvg( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.jit(do_not_specialize=['T']) def chunk_gsa_bwd_k_kernel_intra_dvg( @@ -426,8 +426,8 @@ def chunk_gsa_bwd_k_kernel_intra_dvg( do, dv, dg, - offsets, - indices, + cu_seqlens, + chunk_indices, T, HQ: tl.constexpr, H: tl.constexpr, @@ -444,8 +444,8 @@ def chunk_gsa_bwd_k_kernel_intra_dvg( i_h = i_hq // NG i_t, i_i = i_c // NC, i_c % NC 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -518,7 +518,7 @@ def chunk_gsa_fwd_v( scale: float = 1., initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: _, A, h, ht, o = chunk_gla_fwd( @@ -530,7 +530,7 @@ def chunk_gsa_fwd_v( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return A, h, ht, o @@ -544,7 +544,7 @@ def chunk_gsa_fwd_k( h0: Optional[torch.Tensor] = None, output_final_state: bool = False, scale: float = 1., - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] @@ -553,8 +553,8 @@ def chunk_gsa_fwd_k( BV = min(64, triton.next_power_of_2(V)) HQ = q.shape[2] - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) NG = HQ // H @@ -566,7 +566,7 @@ def chunk_gsa_fwd_k( gv=g, h0=h0, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, states_in_fp32=False ) @@ -580,8 +580,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) g, o, A, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, HQ=HQ, @@ -598,8 +598,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) g, o, A, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, HQ=HQ, H=H, @@ -627,7 +627,7 @@ def chunk_gsa_bwd_v( dht: torch.Tensor, dg: torch.Tensor, scale: float = 1., - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): dq, dk, dv, dg, dh0 = chunk_gla_bwd( @@ -642,7 +642,7 @@ def chunk_gsa_bwd_v( A=A, do=do, dht=dht, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq, dk, dv, dg, dh0 @@ -660,7 +660,7 @@ def chunk_gsa_bwd_k( dht: torch.Tensor, dg: torch.Tensor, scale: float = 1., - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K, V = *k.shape, v.shape[-1] @@ -670,8 +670,8 @@ def chunk_gsa_bwd_k( BV = min(64, triton.next_power_of_2(V)) HQ = q.shape[2] - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) @@ -686,7 +686,7 @@ def chunk_gsa_bwd_k( gv=g, h0=h0, output_final_state=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, states_in_fp32=False ) @@ -701,7 +701,7 @@ def chunk_gsa_bwd_k( h0=h0, dht=dht, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT, states_in_fp32=True ) @@ -712,8 +712,8 @@ def chunk_gsa_bwd_k( g, do, dA, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, B=B, @@ -749,8 +749,8 @@ def chunk_gsa_bwd_k( dg, dgv, dA, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, B=B, @@ -776,8 +776,8 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) do, dv, dg, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, HQ=HQ, H=H, @@ -790,7 +790,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) num_warps=4, num_stages=2 ) - dg = dgv.add_(chunk_local_cumsum(dg, chunk_size=BT, reverse=True, cu_seqlens=offsets)) + dg = dgv.add_(chunk_local_cumsum(dg, chunk_size=BT, reverse=True, cu_seqlens=cu_seqlens)) return dq, dk, dv, dg, dh0 @@ -804,7 +804,7 @@ def chunk_gsa_fwd( initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, output_final_state: bool = False, scale: float = 1., - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: hk0, hv0 = None, None @@ -818,7 +818,7 @@ def chunk_gsa_fwd( h0=hk0, output_final_state=output_final_state, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -834,7 +834,7 @@ def chunk_gsa_fwd( scale=1., initial_state=hv0, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return Ak, hk, hkt, ok, p, Av, hv, hvt, ov @@ -854,7 +854,7 @@ def chunk_gsa_bwd( scale: float, do: torch.Tensor, dht: Tuple[torch.Tensor, torch.Tensor], - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): hk0, hv0 = None, None @@ -878,7 +878,7 @@ def chunk_gsa_bwd( dht=dhvt, dg=None, scale=1., - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -898,7 +898,7 @@ def chunk_gsa_bwd( dht=dhkt, dg=dg, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -925,12 +925,12 @@ def forward( hv0: Optional[torch.Tensor], output_final_state: bool, checkpoint_level: int, - offsets: Optional[torch.LongTensor], + cu_seqlens: Optional[torch.LongTensor], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: T = q.shape[1] chunk_size = min(64, max(16, triton.next_power_of_2(T))) - g_org, g = g, chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=offsets) + g_org, g = g, chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens) Ak, hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( q=q, k=k, @@ -940,7 +940,7 @@ def forward( initial_state=(hk0, hv0), output_final_state=output_final_state, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -957,7 +957,7 @@ def forward( ctx.save_for_backward(q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv) ctx.checkpoint_level = checkpoint_level ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.chunk_size = chunk_size return ov, hkt, hvt @@ -966,11 +966,11 @@ def forward( def backward(ctx, dov, dhkt=None, dhvt=None): q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv = ctx.saved_tensors scale = ctx.scale - offsets = ctx.offsets + cu_seqlens = ctx.cu_seqlens chunk_size = ctx.chunk_size if ctx.checkpoint_level >= 1: - g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=offsets) + g = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens) dq, dk, dv, ds, dg, dhk0, dhv0 = chunk_gsa_bwd( q=q, k=k, @@ -985,7 +985,7 @@ def backward(ctx, dov, dhkt=None, dhvt=None): scale=scale, do=dov, dht=(dhkt, dhvt), - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq, dk, dv, ds, dg, None, dhk0, dhv0, None, None, None, None diff --git a/fla/ops/gsa/fused_recurrent.py b/fla/ops/gsa/fused_recurrent.py index 3b82db7027..47ca8a68a6 100644 --- a/fla/ops/gsa/fused_recurrent.py +++ b/fla/ops/gsa/fused_recurrent.py @@ -151,14 +151,14 @@ def fused_recurrent_gsa_fwd( output_final_state: bool = False, scale: float = 1., reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]: if head_first: B, H, T, K, V, M = *k.shape, v.shape[-1], s.shape[-1] else: B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 HQ = q.shape[1] if head_first else q.shape[2] if HQ != H: raise ValueError("GQA not supported yet.") @@ -186,7 +186,7 @@ def fused_recurrent_gsa_fwd( o=ok, h0=hk0, ht=hkt, - offsets=offsets, + cu_seqlens=cu_seqlens, scale=scale, B=B, T=T, @@ -216,7 +216,7 @@ def fused_recurrent_gsa_fwd( o=ov, h0=hv0, ht=hvt, - offsets=offsets, + cu_seqlens=cu_seqlens, scale=1., B=B, T=T, @@ -249,10 +249,10 @@ def fused_recurrent_gsa_bwd( dhvt: Optional[torch.Tensor] = None, scale: float = 1., reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor]: B, T, H, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV, BM = min(K, 64), min(V, 64), min(M, 64) NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) @@ -279,7 +279,7 @@ def fused_recurrent_gsa_bwd( dv=dv, dht=dhvt, dh0=dhv0, - offsets=offsets, + cu_seqlens=cu_seqlens, scale=1., B=B, T=T, @@ -296,7 +296,7 @@ def fused_recurrent_gsa_bwd( dqv = dqv.sum(0) dsv = dsv.sum(0) dv = dv.sum(0) - dgk = chunk_global_cumsum(dqv * qv.float() - dsv * s.float(), reverse=not reverse, cu_seqlens=offsets) + dgk = chunk_global_cumsum(dqv * qv.float() - dsv * s.float(), reverse=not reverse, cu_seqlens=cu_seqlens) dok = qv * (dqv - (qv * dqv).sum(-1, True)) dq = q.new_empty(NM, B, T, H, K, dtype=torch.float) @@ -318,7 +318,7 @@ def fused_recurrent_gsa_bwd( dv=dsk, dht=dhkt, dh0=dhk0, - offsets=offsets, + cu_seqlens=cu_seqlens, scale=scale, B=B, T=T, @@ -336,7 +336,7 @@ def fused_recurrent_gsa_bwd( dk = dk.sum(0) dsk = dsk.sum(0) - dgv = chunk_global_cumsum(dok.float() * ok.float() - dsk * s.float(), reverse=not reverse, cu_seqlens=offsets) + dgv = chunk_global_cumsum(dok.float() * ok.float() - dsk * s.float(), reverse=not reverse, cu_seqlens=cu_seqlens) ds = dsk.add_(dsv) dg = dgk.add_(dgv) @@ -361,7 +361,7 @@ def forward( hv0: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]: T = q.shape[1] if T == 1 and not q.requires_grad: @@ -386,12 +386,12 @@ def forward( output_final_state=output_final_state, scale=scale, reverse=reverse, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, s, g, qv, hk0, hv0, ok) ctx.scale = scale ctx.reverse = reverse - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return ov.to(q.dtype), hkt, hvt @staticmethod @@ -401,7 +401,7 @@ def backward(ctx, do, dhkt=None, dhvt=None): q, k, v, s, g, qv, hk0, hv0, ok = ctx.saved_tensors scale = ctx.scale reverse = ctx.reverse - offsets = ctx.offsets + cu_seqlens = ctx.cu_seqlens # not supported yet. if dhkt is not None or dhvt is not None: @@ -422,7 +422,7 @@ def backward(ctx, do, dhkt=None, dhvt=None): dhvt=dhvt, scale=scale, reverse=reverse, - offsets=offsets, + cu_seqlens=cu_seqlens, ) return dq.to(q), dk.to(k), dv.to(v), ds.to(s), dg.to(g), None, dhk0, dhv0, None, None, None diff --git a/fla/ops/hgrn/fused_recurrent.py b/fla/ops/hgrn/fused_recurrent.py index d77213c050..e5857482c8 100644 --- a/fla/ops/hgrn/fused_recurrent.py +++ b/fla/ops/hgrn/fused_recurrent.py @@ -14,7 +14,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -31,7 +31,7 @@ def fused_recurrent_hgrn_fwd_kernel( o, h0, ht, - offsets, + cu_seqlens, T, D: tl.constexpr, BD: tl.constexpr, @@ -41,7 +41,7 @@ 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).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 else: bos, eos = i_n * T, i_n * T + T @@ -75,7 +75,7 @@ def fused_recurrent_hgrn_fwd_kernel( @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -95,7 +95,7 @@ def fused_recurrent_hgrn_bwd_kernel( do, dht, dh0, - offsets, + cu_seqlens, T, D: tl.constexpr, BD: tl.constexpr, @@ -105,7 +105,7 @@ 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(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).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 else: bos, eos = i_n * T, i_n * T + T @@ -157,10 +157,10 @@ def fused_recurrent_hgrn_fwd( g: torch.Tensor, initial_state: torch.Tensor = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, D = x.shape - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 o = torch.empty_like(x) final_state = x.new_empty(N, D) if output_final_state else None @@ -172,7 +172,7 @@ def grid(meta): return (triton.cdiv(D, meta['BD']), N) o=o, h0=initial_state, ht=final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, T=T, D=D ) @@ -185,10 +185,10 @@ def fused_recurrent_hgrn_bwd( do: torch.Tensor, dht: torch.Tensor = None, initial_state: torch.Tensor = None, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, D = do.shape - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 dx = torch.empty_like(o, dtype=torch.float) dg = torch.empty_like(g, dtype=torch.float) @@ -203,7 +203,7 @@ def grid(meta): return (triton.cdiv(D, meta['BD']), N) do=do, dht=dht, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, T=T, D=D ) @@ -220,24 +220,24 @@ def forward( g: torch.Tensor, initial_state: torch.Tensor = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ): o, ht = fused_recurrent_hgrn_fwd( x=x, g=g, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets + cu_seqlens=cu_seqlens ) ctx.save_for_backward(g, o, initial_state) - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o, ht @staticmethod @input_guard def backward(ctx, do, dht=None): g, o, initial_state = ctx.saved_tensors - offsets = ctx.offsets + cu_seqlens = ctx.cu_seqlens dx, dg, dh0 = fused_recurrent_hgrn_bwd( g=g, @@ -245,7 +245,7 @@ def backward(ctx, do, dht=None): do=do, dht=dht, initial_state=initial_state, - offsets=offsets + cu_seqlens=cu_seqlens ) return dx, dg, dh0, None, None diff --git a/fla/ops/linear_attn/fused_chunk.py b/fla/ops/linear_attn/fused_chunk.py index 7c0ca03e8a..ea15891bf7 100644 --- a/fla/ops/linear_attn/fused_chunk.py +++ b/fla/ops/linear_attn/fused_chunk.py @@ -46,7 +46,6 @@ def fused_chunk_linear_attn_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, CHECK: tl.constexpr ): - # indices i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H diff --git a/fla/ops/linear_attn/fused_recurrent.py b/fla/ops/linear_attn/fused_recurrent.py index ccf25597eb..7555c64046 100644 --- a/fla/ops/linear_attn/fused_recurrent.py +++ b/fla/ops/linear_attn/fused_recurrent.py @@ -35,7 +35,6 @@ def fused_recurrent_linear_attn_fwd_kernel( USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, ): - # indices i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) p_q = q + i_bh * T*K + i_k * BK + tl.arange(0, BK) diff --git a/fla/ops/nsa/compression.py b/fla/ops/nsa/compression.py index 14444fbdc1..274c932a5d 100644 --- a/fla/ops/nsa/compression.py +++ b/fla/ops/nsa/compression.py @@ -8,13 +8,13 @@ import triton.language as tl from fla.ops.attn.parallel import parallel_attn_bwd_preprocess -from fla.ops.common.utils import prepare_chunk_offsets, prepare_lens, prepare_token_indices +from fla.ops.common.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_token_indices from fla.ops.utils.op import exp, log from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -31,7 +31,7 @@ def parallel_nsa_compression_fwd_kernel( o, lse, scale, - offsets, + cu_seqlens, token_indices, chunk_offsets, T, @@ -51,7 +51,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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) else: @@ -116,7 +116,7 @@ def parallel_nsa_compression_fwd_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -135,7 +135,7 @@ def parallel_nsa_compression_bwd_kernel_dq( do, dq, scale, - offsets, + cu_seqlens, token_indices, chunk_offsets, T, @@ -156,7 +156,7 @@ def parallel_nsa_compression_bwd_kernel_dq( 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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) else: @@ -218,7 +218,7 @@ def parallel_nsa_compression_bwd_kernel_dq( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -237,7 +237,7 @@ def parallel_nsa_compression_bwd_kernel_dkv( do, dk, dv, - offsets, + cu_seqlens, chunk_indices, chunk_offsets, scale, @@ -259,7 +259,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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) else: @@ -320,7 +320,7 @@ def parallel_nsa_compression_fwd( v: torch.Tensor, block_size: int, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, token_indices: Optional[torch.LongTensor] = None, ): B, T, HQ, K, V = *q.shape, v.shape[-1] @@ -337,7 +337,7 @@ def parallel_nsa_compression_fwd( NV = triton.cdiv(V, BV) assert NK == 1, "The key dimension can not be larger than 256" - chunk_offsets = prepare_chunk_offsets(offsets, BS) if offsets is not None else None + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) if cu_seqlens is not None else None grid = (T, NV, B * H) o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) @@ -350,7 +350,7 @@ def parallel_nsa_compression_fwd( o=o, lse=lse, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices, chunk_offsets=chunk_offsets, T=T, @@ -376,7 +376,7 @@ def parallel_nsa_compression_bwd( do: torch.Tensor, block_size: int = 64, scale: float = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, token_indices: Optional[torch.LongTensor] = None, ): B, T, HQ, K, V = *q.shape, v.shape[-1] @@ -386,11 +386,8 @@ def parallel_nsa_compression_bwd( BK = triton.next_power_of_2(K) BV = min(128, triton.next_power_of_2(v.shape[-1])) NV = triton.cdiv(V, BV) - if offsets is not None: - lens = prepare_lens(offsets) - chunk_indices = torch.cat([torch.arange(n) for n in triton.cdiv(triton.cdiv(lens, BS), BC).tolist()]) - chunk_indices = torch.stack([chunk_indices.eq(0).cumsum(0) - 1, chunk_indices], 1).to(offsets) - chunk_offsets = prepare_chunk_offsets(offsets, BS) + if cu_seqlens is not None: + chunk_indices, chunk_offsets = prepare_chunk_indices(cu_seqlens, BS), prepare_chunk_offsets(cu_seqlens, BS) NC = len(chunk_indices) else: chunk_indices, chunk_offsets = None, None @@ -409,7 +406,7 @@ def parallel_nsa_compression_bwd( do=do, dq=dq, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices, chunk_offsets=chunk_offsets, T=T, @@ -439,7 +436,7 @@ def parallel_nsa_compression_bwd( do=do, dk=dk, dv=dv, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_offsets=chunk_offsets, scale=scale, @@ -471,15 +468,15 @@ def forward( v, block_size, scale, - offsets + cu_seqlens ): ctx.dtype = q.dtype - # 2-d sequence indices denoting the offsets of tokens in each sequence - # for example, if the passed `offsets` is [0, 2, 6], + # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence + # for example, if the passed `cu_seqlens` is [0, 2, 6], # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] - token_indices = prepare_token_indices(offsets) if offsets is not None else None + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None o, lse = parallel_nsa_compression_fwd( q=q, @@ -487,11 +484,11 @@ def forward( v=v, block_size=block_size, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices ) ctx.save_for_backward(q, k, v, o, lse) - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.token_indices = token_indices ctx.block_size = block_size ctx.scale = scale @@ -511,7 +508,7 @@ def backward(ctx, do, *args): do=do, block_size=ctx.block_size, scale=ctx.scale, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, token_indices=ctx.token_indices ) return dq.to(q), dk.to(k), dv.to(v), None, None, None @@ -523,7 +520,7 @@ def parallel_nsa_compression( v: torch.Tensor, block_size: int = 64, scale: float = None, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ): if scale is None: scale = k.shape[-1] ** -0.5 @@ -533,5 +530,5 @@ def parallel_nsa_compression( v, block_size, scale, - offsets + cu_seqlens ) diff --git a/fla/ops/nsa/naive.py b/fla/ops/nsa/naive.py index c365112e3e..fc95e9d05b 100644 --- a/fla/ops/nsa/naive.py +++ b/fla/ops/nsa/naive.py @@ -11,7 +11,7 @@ def naive_nsa( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - indices: torch.LongTensor, + block_indices: torch.LongTensor, block_size: int = 64, scale: Optional[float] = None, cu_seqlens: Optional[torch.LongTensor] = None, @@ -26,7 +26,7 @@ def naive_nsa( GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16. v (torch.Tensor): values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. - indices (torch.LongTensor): + block_indices (torch.LongTensor): Block indices of shape `[B, T, H, S]` if `head_first=False` else `[B, H, T, S]`. `S` is the number of selected blocks for each query token, which is set to 16 in the paper. block_size (int): @@ -52,12 +52,12 @@ def naive_nsa( "Sequences with variable lengths are not supported for head-first mode" ) if head_first: - q, k, v, indices = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v, indices)) + q, k, v, block_indices = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v, block_indices)) dtype = q.dtype G = q.shape[2] // k.shape[2] BS = block_size - k, v, indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, indices)) + k, v, block_indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, block_indices)) q, k, v = map(lambda x: x.float(), (q, k, v)) o = torch.zeros_like(v) @@ -65,18 +65,20 @@ def naive_nsa( if cu_seqlens is None: varlen = False B, T = q.shape[:2] - cu_seqlens = torch.cat([indices.new_tensor(range(0, B*T, T)), indices.new_tensor([B*T])]) + cu_seqlens = torch.cat([ + block_indices.new_tensor(range(0, B*T, T)), block_indices.new_tensor([B*T]) + ]) for i in range(len(cu_seqlens) - 1): if not varlen: - q_b, k_b, v_b, i_b = q[i], k[i], v[i], indices[i] + q_b, k_b, v_b, i_b = q[i], k[i], v[i], block_indices[i] else: T = cu_seqlens[i+1] - cu_seqlens[i] - q_b, k_b, v_b, i_b = map(lambda x: x[0][cu_seqlens[i]:cu_seqlens[i+1]], (q, k, v, indices)) + q_b, k_b, v_b, i_b = map(lambda x: x[0][cu_seqlens[i]:cu_seqlens[i+1]], (q, k, v, block_indices)) i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS)) # [T, S*BS, HQ] - i_b = i_b.view(T, indices.shape[2], -1).transpose(1, 2) + i_b = i_b.view(T, block_indices.shape[2], -1).transpose(1, 2) for i_q in range(T): # [HQ, D] q_i = q_b[i_q] * scale diff --git a/fla/ops/nsa/parallel.py b/fla/ops/nsa/parallel.py index 98cdee2a17..6784bdae4f 100644 --- a/fla/ops/nsa/parallel.py +++ b/fla/ops/nsa/parallel.py @@ -28,7 +28,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -44,7 +44,7 @@ def parallel_nsa_kernel_topk( lse, scale, block_indices, - offsets, + cu_seqlens, token_indices, chunk_offsets, T, @@ -63,7 +63,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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) else: @@ -160,7 +160,7 @@ def parallel_nsa_kernel_topk( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), }) @triton.autotune( @@ -180,7 +180,7 @@ def parallel_nsa_fwd_kernel( scale, block_indices, block_counts, - offsets, + cu_seqlens, token_indices, T, H: tl.constexpr, @@ -200,7 +200,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 else: bos, eos = i_b * T, i_b * T + T @@ -286,7 +286,7 @@ def parallel_nsa_kernel_mask( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor) }) @triton.autotune( @@ -308,7 +308,7 @@ def parallel_nsa_bwd_kernel_dq( scale, block_indices, block_counts, - offsets, + cu_seqlens, token_indices, T, B: tl.constexpr, @@ -329,7 +329,7 @@ def parallel_nsa_bwd_kernel_dq( 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 else: bos, eos = i_b * T, i_b * T + T @@ -394,7 +394,7 @@ def parallel_nsa_bwd_kernel_dq( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -414,7 +414,7 @@ def parallel_nsa_bwd_kernel_dkv( dk, dv, block_mask, - offsets, + cu_seqlens, chunk_indices, scale, T, @@ -435,7 +435,7 @@ def parallel_nsa_bwd_kernel_dkv( 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 else: bos, eos = i_b * T, i_b * T + T @@ -492,7 +492,7 @@ def parallel_nsa_topk( block_counts: Union[torch.LongTensor, int], block_size: int = 64, scale: float = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ) -> torch.LongTensor: B, T, HQ, K = q.shape H = k.shape[2] @@ -504,8 +504,8 @@ def parallel_nsa_topk( BK = triton.next_power_of_2(K) block_indices = torch.zeros(B, T, H, S, dtype=torch.int32, device=q.device) - token_indices = prepare_token_indices(offsets) if offsets is not None else None - chunk_offsets = prepare_chunk_offsets(offsets, BS) if offsets is not None else None + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) if cu_seqlens is not None else None grid = (T, B * H) parallel_nsa_kernel_topk[grid]( q=q, @@ -513,7 +513,7 @@ def parallel_nsa_topk( lse=lse, scale=scale, block_indices=block_indices, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices, chunk_offsets=chunk_offsets, T=T, @@ -537,7 +537,7 @@ def parallel_nsa_fwd( block_counts: Union[torch.LongTensor, int], block_size: int, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, token_indices: Optional[torch.LongTensor] = None, ): B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] @@ -567,7 +567,7 @@ def parallel_nsa_fwd( scale=scale, block_indices=block_indices, block_counts=block_counts, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices, T=T, H=H, @@ -586,13 +586,13 @@ def parallel_nsa_fwd( def parallel_nsa_block_mask( block_indices: torch.LongTensor, block_counts: Union[torch.LongTensor, int], - offsets: torch.LongTensor, + cu_seqlens: torch.LongTensor, block_size: int, ): B, T, H, S = block_indices.shape BS = block_size - if offsets is not None: - NS = triton.cdiv(prepare_lens(offsets).max().item(), BS) + if cu_seqlens is not None: + NS = triton.cdiv(prepare_lens(cu_seqlens).max().item(), BS) else: NS = triton.cdiv(T, BS) block_mask = torch.zeros(B, T, H, NS, dtype=torch.bool, device=block_indices.device) @@ -621,7 +621,7 @@ def parallel_nsa_bwd( block_counts: Union[torch.LongTensor, int], block_size: int = 64, scale: float = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, token_indices: Optional[torch.LongTensor] = None, ): B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] @@ -646,7 +646,7 @@ def parallel_nsa_bwd( dq=dq, block_indices=block_indices, block_counts=block_counts, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices, scale=scale, T=T, @@ -663,15 +663,15 @@ def parallel_nsa_bwd( ) dq = dq.sum(0) - if offsets is not None: - chunk_indices = prepare_chunk_indices(offsets, BS) + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BS) NS = len(chunk_indices) else: chunk_indices = None NS = triton.cdiv(T, BS) # [B, T, H, M] - block_mask = parallel_nsa_block_mask(block_indices, block_counts, offsets, block_size) + block_mask = parallel_nsa_block_mask(block_indices, block_counts, cu_seqlens, block_size) dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) @@ -686,7 +686,7 @@ def parallel_nsa_bwd( dk=dk, dv=dv, block_mask=block_mask, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, scale=scale, T=T, @@ -711,14 +711,14 @@ class ParallelNSAFunction(torch.autograd.Function): @staticmethod @contiguous @autocast_custom_fwd - def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, offsets): + def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens): ctx.dtype = q.dtype - # 2-d sequence indices denoting the offsets of tokens in each sequence - # for example, if the passed `offsets` is [0, 2, 6], + # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence + # for example, if the passed `cu_seqlens` is [0, 2, 6], # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] - token_indices = prepare_token_indices(offsets) if offsets is not None else None + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None o, lse = parallel_nsa_fwd( q=q, @@ -728,13 +728,13 @@ def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, offset block_counts=block_counts, block_size=block_size, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, token_indices=token_indices ) ctx.save_for_backward(q, k, v, o, lse) ctx.block_indices = block_indices ctx.block_counts = block_counts - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens ctx.token_indices = token_indices ctx.block_size = block_size ctx.scale = scale @@ -756,7 +756,7 @@ def backward(ctx, do): block_counts=ctx.block_counts, block_size=ctx.block_size, scale=ctx.scale, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, token_indices=ctx.token_indices ) return dq.to(q), dk.to(k), dv.to(v), None, None, None, None, None, None, None, None @@ -766,9 +766,9 @@ def parallel_nsa( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - g_cmp: torch.Tensor, - g_slc: torch.Tensor, - g_swa: torch.Tensor, + g_cmp: Optional[torch.Tensor] = None, + g_slc: Optional[torch.Tensor] = None, + g_swa: Optional[torch.Tensor] = None, block_indices: Optional[torch.LongTensor] = None, block_counts: Union[torch.LongTensor, int] = 16, block_size: int = 64, @@ -839,7 +839,7 @@ def parallel_nsa( v=v_cmp, block_size=block_size, scale=scale, - offsets=cu_seqlens + cu_seqlens=cu_seqlens ) if block_indices is not None: warnings.warn("`block_indices` will be ignored when `g_cmp` is provided") @@ -850,10 +850,11 @@ def parallel_nsa( block_counts=block_counts, block_size=block_size, scale=scale, - offsets=cu_seqlens + cu_seqlens=cu_seqlens ) - o_slc = ParallelNSAFunction.apply(q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens) - o = o_slc * g_slc.unsqueeze(-1) + o = o_slc = ParallelNSAFunction.apply(q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens) + if g_slc is not None: + o = o_slc * g_slc.unsqueeze(-1) if o_cmp is not None: o = torch.addcmul(o, o_cmp, g_cmp.unsqueeze(-1)) if window_size > 0: diff --git a/fla/ops/retention/fused_chunk.py b/fla/ops/retention/fused_chunk.py index 6089634b1b..cbb56bbe8e 100644 --- a/fla/ops/retention/fused_chunk.py +++ b/fla/ops/retention/fused_chunk.py @@ -32,7 +32,6 @@ def fused_chunk_retention_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, CHECK: tl.constexpr ): - # indices i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_h = i_bh % H diff --git a/fla/ops/rwkv6/chunk.py b/fla/ops/rwkv6/chunk.py index fde191beb9..25b24b0ee1 100644 --- a/fla/ops/rwkv6/chunk.py +++ b/fla/ops/rwkv6/chunk.py @@ -10,7 +10,7 @@ from einops import rearrange from fla.ops.common.chunk_h import chunk_fwd_h -from fla.ops.common.utils import prepare_chunk_indices +from fla.ops.common.utils import prepare_chunk_indices, prepare_chunk_offsets from fla.ops.gla.chunk import chunk_gla_bwd_dA, chunk_gla_bwd_dv, chunk_gla_fwd_o_gk from fla.ops.utils.op import exp from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, input_guard, use_cuda_graph @@ -20,7 +20,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -37,8 +37,8 @@ def chunk_rwkv6_fwd_cumsum_kernel( s, oi, oe, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, S: tl.constexpr, @@ -49,8 +49,8 @@ def chunk_rwkv6_fwd_cumsum_kernel( i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -73,12 +73,12 @@ def chunk_rwkv6_fwd_cumsum_kernel( def chunk_rwkv6_fwd_cumsum( g: torch.Tensor, chunk_size: int, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, ) -> torch.Tensor: B, T, H, S = g.shape BT = chunk_size - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) gi, ge = torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float) def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) @@ -87,8 +87,8 @@ def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) g, gi, ge, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, S=S, @@ -98,7 +98,7 @@ def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -117,8 +117,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( gi, # cumulative decay inclusive ge, # cumulative decay exclusive A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -133,8 +133,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -175,14 +175,12 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - triton.Config({}, num_warps=8), + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] ], key=['BK', 'BT'], use_cuda_graph=use_cuda_graph, @@ -195,8 +193,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( ge, u, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -210,8 +208,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( i_b, i_h = i_bh // H, i_bh % H i_j = i_i 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -249,7 +247,7 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -269,8 +267,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( ge, u, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, B: tl.constexpr, T, @@ -287,8 +285,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( i_t, i_i = i_tc // NC, i_tc % NC i_j = i_i 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) + 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) all = T T = eos - bos else: @@ -329,7 +327,7 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -345,8 +343,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( A, A2, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -358,8 +356,8 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H 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) + 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) all = T T = eos - bos else: @@ -380,7 +378,7 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -402,7 +400,7 @@ def chunk_rwkv6_bwd_kernel_dh( dh, dht, dh0, - offsets, + cu_seqlens, chunk_offsets, scale, T, @@ -422,7 +420,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -463,7 +461,7 @@ def chunk_rwkv6_bwd_kernel_dh( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -482,8 +480,8 @@ def chunk_rwkv6_bwd_kernel_intra( dA, dq, dk, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, K: tl.constexpr, @@ -497,8 +495,8 @@ def chunk_rwkv6_bwd_kernel_intra( i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC 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) + 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) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -601,7 +599,7 @@ def chunk_rwkv6_bwd_kernel_intra( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -631,8 +629,8 @@ def chunk_rwkv6_bwd_kernel_inter( dk2, dg, du, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -648,8 +646,8 @@ def chunk_rwkv6_bwd_kernel_inter( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -731,14 +729,14 @@ def chunk_rwkv6_fwd_intra( ge: torch.Tensor, u: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K = k.shape BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(16, BT) NC = triton.cdiv(BT, BC) @@ -750,8 +748,8 @@ def chunk_rwkv6_fwd_intra( gi, ge, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -772,8 +770,8 @@ def chunk_rwkv6_fwd_intra( ge, u, A, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -796,8 +794,8 @@ def chunk_rwkv6_fwd_intra( ge, u, A_intra, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, B=B, T=T, @@ -813,8 +811,8 @@ def chunk_rwkv6_fwd_intra( chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge[grid]( A_intra, A, - offsets, - indices, + cu_seqlens, + chunk_indices, B=B, T=T, H=H, @@ -835,7 +833,7 @@ def chunk_rwkv6_bwd_dh( h0: torch.Tensor, dht: torch.Tensor, scale: float, - offsets: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, chunk_size: int = 64, states_in_fp32: bool = False ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -844,12 +842,12 @@ def chunk_rwkv6_bwd_dh( BT = min(chunk_size, max(16, triton.next_power_of_2(T))) # N: the actual number of sequences in the batch with either equal or variable lengths # NG: number of groups in GQA - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - if offsets is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT = len(offsets) - 1, len(indices) - chunk_offsets = torch.cat([offsets.new_tensor([0]), triton.cdiv(offsets[1:] - offsets[:-1], BT)]).cumsum(-1) + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) NG = HQ // H dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) @@ -864,7 +862,7 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), dh=dh, dht=dht, dh0=dh0, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, scale=scale, T=T, @@ -884,7 +882,7 @@ def chunk_rwkv6_bwd_dqk_intra( gi: torch.Tensor, ge: torch.Tensor, dA: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K = q.shape @@ -892,8 +890,8 @@ def chunk_rwkv6_bwd_dqk_intra( BC = min(16, BT) BK = min(64, triton.next_power_of_2(K)) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) @@ -908,8 +906,8 @@ def chunk_rwkv6_bwd_dqk_intra( dA, dq, dk, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, K=K, @@ -936,14 +934,14 @@ def chunk_rwkv6_bwd_dqkgu( dq: torch.Tensor, dk: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): B, T, H, K, V = *k.shape, v.shape[-1] BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dq2 = torch.empty_like(dq) dk2 = torch.empty_like(dk) @@ -967,8 +965,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) dk2, dg, du, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -989,10 +987,10 @@ def chunk_rwkv6_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, offsets=offsets) + gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) h, ht = chunk_fwd_h( k=k, v=v, @@ -1001,7 +999,7 @@ def chunk_rwkv6_fwd( gv=None, h0=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, states_in_fp32=True ) @@ -1014,7 +1012,7 @@ def chunk_rwkv6_fwd( ge=ge, u=u, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -1025,7 +1023,7 @@ def chunk_rwkv6_fwd( A=A, h=h, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return A, h, ht, o @@ -1042,10 +1040,10 @@ def chunk_rwkv6_bwd( A: torch.Tensor, do: torch.Tensor, dht: torch.Tensor, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ): - gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, offsets=offsets) + gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) h, _ = chunk_fwd_h( k=k, v=v, @@ -1054,7 +1052,7 @@ def chunk_rwkv6_bwd( gv=None, h0=initial_state, output_final_state=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, states_in_fp32=True ) @@ -1068,7 +1066,7 @@ def chunk_rwkv6_bwd( h0=initial_state, dht=dht, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, states_in_fp32=True ) @@ -1078,7 +1076,7 @@ def chunk_rwkv6_bwd( v=v, do=do, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dv = chunk_gla_bwd_dv( @@ -1087,7 +1085,7 @@ def chunk_rwkv6_bwd( A=A, do=do, dh=dh, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dq, dk = chunk_rwkv6_bwd_dqk_intra( @@ -1096,7 +1094,7 @@ def chunk_rwkv6_bwd( gi=gi, ge=ge, dA=dA, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dq, dk, dg, du = chunk_rwkv6_bwd_dqkgu( @@ -1114,7 +1112,7 @@ def chunk_rwkv6_bwd( dq=dq, dk=dk, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq, dk, dv, dg, du, dh0 @@ -1135,7 +1133,7 @@ def forward( scale, initial_state, output_final_state, - offsets, + cu_seqlens, ): T = q.shape[1] if check_shared_mem(): @@ -1152,7 +1150,7 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) @@ -1160,7 +1158,7 @@ def forward( ctx.chunk_size = chunk_size ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o, ht @staticmethod @@ -1168,7 +1166,7 @@ def forward( @autocast_custom_bwd def backward(ctx, do, dht): q, k, v, g, initial_state, A, u = ctx.saved_tensors - chunk_size, scale, offsets = ctx.chunk_size, ctx.scale, ctx.offsets + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens dq, dk, dv, dg, du, dh0 = chunk_rwkv6_bwd( q=q, k=k, @@ -1180,7 +1178,7 @@ def backward(ctx, do, dht): A=A, do=do, dht=dht, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq.to(q), dk.to(k), dv.to(v), dg.to(g), du.to(u), None, dh0, None, None @@ -1188,10 +1186,10 @@ def backward(ctx, do, dht): @torch.compiler.disable def chunk_rwkv6( - q: torch.Tensor, + r: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - g: torch.Tensor, + w: torch.Tensor, u: torch.Tensor, scale: Optional[int] = None, initial_state: torch.Tensor = None, @@ -1201,13 +1199,13 @@ def chunk_rwkv6( ) -> Tuple[torch.Tensor, torch.Tensor]: r""" Args: - q (torch.Tensor): + r (torch.Tensor): queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. k (torch.Tensor): keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. v (torch.Tensor): values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. - g (torch.Tensor): + w (torch.Tensor): Forget gates of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]` applied to keys. u (torch.Tensor): bonus representations of shape `[H]`. @@ -1240,23 +1238,23 @@ def chunk_rwkv6( >>> from fla.ops.rwkv6 import chunk_rwkv6 # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') + >>> r = torch.randn(B, T, H, K, device='cuda') >>> k = torch.randn(B, T, H, K, device='cuda') >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> w = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) >>> u = torch.randn(H, K, device='cuda') >>> h0 = torch.randn(B, H, K, V, device='cuda') >>> o, ht = chunk_rwkv6( - q, k, v, g, u, + r, k, v, w, u, initial_state=h0, output_final_state=True ) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + >>> r, k, v, w = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (r, k, v, w)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> cu_seqlens = r.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) >>> o_var, ht_var = chunk_rwkv6( - q, k, v, g, u, + r, k, v, w, u, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens @@ -1269,18 +1267,18 @@ def chunk_rwkv6( "head_first is deprecated and will be removed in a future version. " "Please use head_first=False for now instead." ) - q, k, v, g = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v, g)) - if not head_first and q.shape[1] < q.shape[2]: + r, k, v, w = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (r, k, v, w)) + if not head_first and r.shape[1] < r.shape[2]: warnings.warn( - f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " "This may indicate the inputs were passed in head-first format [B, H, T, ...] " "when head_first=False was specified. " "Please verify your input tensor format matches the expected shape [B, T, H, ...]." ) if cu_seqlens is not None: - if q.shape[0] != 1: + if r.shape[0] != 1: raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`." f"Please flatten variable-length inputs before processing." ) if head_first: @@ -1293,12 +1291,12 @@ def chunk_rwkv6( f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." ) if scale is None: - scale = q.shape[-1] ** -0.5 + scale = r.shape[-1] ** -0.5 o, final_state = ChunkRWKV6Function.apply( - q, + r, k, v, - g, + w, u, scale, initial_state, diff --git a/fla/ops/rwkv6/fused_recurrent.py b/fla/ops/rwkv6/fused_recurrent.py index 9b006f5030..8b5463e791 100644 --- a/fla/ops/rwkv6/fused_recurrent.py +++ b/fla/ops/rwkv6/fused_recurrent.py @@ -16,7 +16,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -35,7 +35,7 @@ def fused_recurrent_rwkv6_fwd_kernel( o, # output [NK, B, H, T, V]/[NK, B, T, H, V] h0, # initial hidden state [B, H, K, V] ht, # final hidden state [B, H, K, V] - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -52,7 +52,7 @@ def fused_recurrent_rwkv6_fwd_kernel( i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -101,7 +101,7 @@ def fused_recurrent_rwkv6_fwd_kernel( @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -121,7 +121,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dq( dq, # gradient of query [NV, B, H, T, K]/[NV, B, T, H, K] dq1, # gradient of query_aux [NV, B, H, T, K]/[NV, B, T, H, K] h0, - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -137,7 +137,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dq( i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -190,7 +190,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dq( @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -212,7 +212,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dkv( dk1, # gradient of key_aux [NV, B, H, T, K]/[NK, B, T, H, K] dv, # gradient of value [NK, B, H, T, V]/[NV, B, T, H, V] dh0, # gradient of initial hidden state [N, H, K, V] - offsets, + cu_seqlens, scale, T, B: tl.constexpr, @@ -228,7 +228,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dkv( i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + 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 else: @@ -286,7 +286,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dkv( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -304,7 +304,7 @@ def fused_recurrent_rwkv6_bwd_kernel_dw( dq, dk, dw, - offsets, + cu_seqlens, scale, T, H: tl.constexpr, @@ -317,7 +317,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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) else: bos, eos = i_n * T, i_n * T + T T = eos - bos @@ -359,10 +359,10 @@ def fused_recurrent_rwkv6_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) @@ -380,7 +380,7 @@ def fused_recurrent_rwkv6_fwd( o, h0, ht, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -405,10 +405,10 @@ def fused_recurrent_rwkv6_bwd( scale: Optional[float] = None, initial_state: Optional[torch.Tensor] = None, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = min(triton.next_power_of_2(K), 16), min(triton.next_power_of_2(V), 64) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) @@ -426,7 +426,7 @@ def fused_recurrent_rwkv6_bwd( dq, dq1, initial_state, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -460,7 +460,7 @@ def fused_recurrent_rwkv6_bwd( dk1, dv, dh0, - offsets, + cu_seqlens, scale, T=T, B=B, @@ -483,7 +483,7 @@ def grid(meta): return (triton.cdiv(meta['K'], meta['BK']), N * H) dq1, dk1, dw, - offsets, + cu_seqlens, scale, T=T, H=H, @@ -511,7 +511,7 @@ def forward( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, reverse: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): o, ht = fused_recurrent_rwkv6_fwd( q=q, @@ -523,12 +523,12 @@ def forward( initial_state=initial_state, output_final_state=output_final_state, reverse=reverse, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, w, u, initial_state) ctx.scale = scale ctx.reverse = reverse - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o.to(v), ht @staticmethod @@ -547,7 +547,7 @@ def backward(ctx, do, dht): scale=ctx.scale, initial_state=initial_state, reverse=ctx.reverse, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, ) dh0 = dh0.to(initial_state) if dh0 is not None else dh0 return dq.to(q), dk.to(k), dv.to(v), dw.to(w), du.to(u), None, dh0, None, None, None diff --git a/fla/ops/simple_gla/chunk.py b/fla/ops/simple_gla/chunk.py index bfc94e9ecd..324d050809 100644 --- a/fla/ops/simple_gla/chunk.py +++ b/fla/ops/simple_gla/chunk.py @@ -22,10 +22,10 @@ def chunk_simple_gla_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=offsets) if g is not None else None + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) if g is not None else None h, ht = chunk_fwd_h( k=k, v=v, @@ -35,7 +35,7 @@ def chunk_simple_gla_fwd( h0=initial_state, output_final_state=output_final_state, states_in_fp32=False, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) o = chunk_fwd_o( @@ -45,7 +45,7 @@ def chunk_simple_gla_fwd( g=g, h=h, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return g, o, ht @@ -60,7 +60,7 @@ def chunk_simple_gla_bwd( do: torch.Tensor, dht: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # (SY 09/22) states_in_fp32 seems not affecting the error of dg but for safety, set to True @@ -73,7 +73,7 @@ def chunk_simple_gla_bwd( h0=initial_state, output_final_state=False, states_in_fp32=True, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dh, dh0 = chunk_bwd_dh( @@ -88,7 +88,7 @@ def chunk_simple_gla_bwd( dht=dht, scale=scale, states_in_fp32=True, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dq, dk, _, dg = chunk_bwd_dqkwg( @@ -100,7 +100,7 @@ def chunk_simple_gla_bwd( do=do, dh=dh, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) dv = chunk_bwd_dv( @@ -110,7 +110,7 @@ def chunk_simple_gla_bwd( do=do, dh=dh, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) return dq, dk, dv, dg, dh0 @@ -130,7 +130,7 @@ def forward( scale, initial_state, output_final_state, - offsets + cu_seqlens ): T = q.shape[1] chunk_size = min(64, max(16, triton.next_power_of_2(T))) @@ -143,20 +143,20 @@ def forward( scale=scale, initial_state=initial_state, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) ctx.save_for_backward(q, k, v, g, initial_state) ctx.chunk_size = chunk_size ctx.scale = scale - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o.to(q.dtype), ht @staticmethod @input_guard @autocast_custom_bwd def backward(ctx, do, dht): - chunk_size, scale, offsets = ctx.chunk_size, ctx.scale, ctx.offsets + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens q, k, v, g, initial_state = ctx.saved_tensors dq, dk, dv, dg, dh0 = chunk_simple_gla_bwd( q=q, @@ -167,12 +167,15 @@ def backward(ctx, do, dht): do=do, dht=dht, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=chunk_size ) - dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=offsets) if g is not None else None + if g is not None: + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens).to(g) + else: + dg = None - return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dg.to(g.dtype), None, dh0, None, None + return dq.to(q), dk.to(k), dv.to(v), dg, None, dh0, None, None @torch.compiler.disable diff --git a/fla/ops/simple_gla/parallel.py b/fla/ops/simple_gla/parallel.py index ed431f8fdf..37228da5d1 100644 --- a/fla/ops/simple_gla/parallel.py +++ b/fla/ops/simple_gla/parallel.py @@ -29,8 +29,8 @@ @triton.heuristics({ 'NV': lambda args: triton.cdiv(args['V'], args['BV']), 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_G': lambda args: args['g'] is not None + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -49,8 +49,8 @@ def parallel_simple_gla_fwd_kernel( o, attn, scale, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -71,8 +71,8 @@ def parallel_simple_gla_fwd_kernel( o += i_k * B * T * H * V 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -364,8 +364,8 @@ def parallel_simple_gla_bwd_kernel_dkv( @triton.heuristics({ 'NV': lambda args: triton.cdiv(args['V'], args['BV']), - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_G': lambda args: args['g'] is not None + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -386,8 +386,8 @@ def parallel_simple_gla_bwd_kernel( dv, dg, scale, - offsets, - indices, + cu_seqlens, + chunk_indices, T, B: tl.constexpr, H: tl.constexpr, @@ -411,8 +411,8 @@ def parallel_simple_gla_bwd_kernel( dg += i_kv * B * H * T 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -491,7 +491,7 @@ def parallel_simple_gla_fwd( scale: float, output_attentions: bool = False, chunk_size: int = 128, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] BT, BS = chunk_size, 32 @@ -509,12 +509,12 @@ def parallel_simple_gla_fwd( NV = triton.cdiv(V, BV) assert BT % BS == 0 - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) # local cumulative decay in log space if g is not None: - g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=offsets) + g = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens) grid = (NK * NV, NT, B * H) o = torch.empty(NK, *v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) attn = q.new_zeros(NK, B, H, T, T) if output_attentions else None @@ -527,8 +527,8 @@ def parallel_simple_gla_fwd( o=o, attn=attn, scale=scale, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, B=B, H=H, T=T, @@ -554,7 +554,7 @@ def parallel_simple_gla_bwd( do: torch.Tensor, scale: float, chunk_size: int = 128, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): B, T, H, K, V = *k.shape, v.shape[-1] BT, BS = chunk_size, 32 @@ -580,8 +580,8 @@ def parallel_simple_gla_bwd( dv = torch.empty(NK, * v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) dg = torch.empty(NK*NV, *g.shape, dtype=torch.float, device=q.device) if g is not None else None - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) grid = (NK * NV, NT, B * H) parallel_simple_gla_bwd_kernel[grid]( @@ -594,7 +594,8 @@ def parallel_simple_gla_bwd( dk=dk, dv=dv, dg=dg, - offsets=offsets, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, T=T, B=B, @@ -609,7 +610,7 @@ def parallel_simple_gla_bwd( dq = dq.sum(0) dk = dk.sum(0) dv = dv.sum(0) - dg = chunk_global_cumsum(dg.sum(0), reverse=True, cu_seqlens=offsets) if g is not None else None + dg = chunk_global_cumsum(dg.sum(0), reverse=True, cu_seqlens=cu_seqlens) if g is not None else None return dq, dk, dv, dg @@ -618,7 +619,7 @@ class ParallelSimpleGLAFunction(torch.autograd.Function): @staticmethod @input_guard @autocast_custom_fwd - def forward(ctx, q, k, v, g, scale, output_attentions, offsets): + def forward(ctx, q, k, v, g, scale, output_attentions, cu_seqlens): chunk_size = 128 ctx.dtype = q.dtype @@ -630,9 +631,9 @@ def forward(ctx, q, k, v, g, scale, output_attentions, offsets): scale=scale, output_attentions=output_attentions, chunk_size=chunk_size, - offsets=offsets, + cu_seqlens=cu_seqlens, ) - ctx.save_for_backward(q, k, v, g, offsets) + ctx.save_for_backward(q, k, v, g, cu_seqlens) ctx.scale = scale ctx.chunk_size = chunk_size return o.to(q.dtype), attn @@ -641,7 +642,7 @@ def forward(ctx, q, k, v, g, scale, output_attentions, offsets): @input_guard @autocast_custom_bwd def backward(ctx, do, da=None): - q, k, v, g, offsets = ctx.saved_tensors + q, k, v, g, cu_seqlens = ctx.saved_tensors dq, dk, dv, dg = parallel_simple_gla_bwd( q=q, k=k, @@ -650,7 +651,7 @@ def backward(ctx, do, da=None): do=do, scale=ctx.scale, chunk_size=ctx.chunk_size, - offsets=offsets, + cu_seqlens=cu_seqlens, ) return dq.to(q), dk.to(k), dv.to(v), dg.to(ctx.dtype) if dg is not None else None, None, None, None diff --git a/fla/ops/ttt/chunk.py b/fla/ops/ttt/chunk.py index dd6970be5d..bb8c1d6360 100755 --- a/fla/ops/ttt/chunk.py +++ b/fla/ops/ttt/chunk.py @@ -19,7 +19,7 @@ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -43,7 +43,7 @@ def chunk_ttt_linear_fwd_kernel_h( hb0, ht, hbt, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -61,7 +61,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -123,7 +123,7 @@ def chunk_ttt_linear_fwd_kernel_h( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -142,8 +142,8 @@ def chunk_ttt_linear_fwd_kernel_o( h, hb, o, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -159,8 +159,8 @@ def chunk_ttt_linear_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -216,7 +216,7 @@ def chunk_ttt_linear_fwd_kernel_o( @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -240,7 +240,7 @@ def chunk_ttt_linear_bwd_kernel_h( x, y, r, - offsets, + cu_seqlens, chunk_offsets, T, H: tl.constexpr, @@ -257,7 +257,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -317,7 +317,7 @@ def chunk_ttt_linear_bwd_kernel_h( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -333,8 +333,8 @@ def chunk_ttt_linear_bwd_kernel_dv_local( eta, do, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, H: tl.constexpr, @@ -348,8 +348,8 @@ def chunk_ttt_linear_bwd_kernel_dv_local( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -391,7 +391,7 @@ def chunk_ttt_linear_bwd_kernel_dv_local( 'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None, 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, 'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -425,7 +425,7 @@ def chunk_ttt_linear_bwd_kernel_norm( dk, dw, db, - offsets, + cu_seqlens, chunk_offsets, scale, T, @@ -444,7 +444,7 @@ 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(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -547,7 +547,7 @@ def chunk_ttt_linear_bwd_kernel_norm( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -570,8 +570,8 @@ def chunk_bwd_kernel_dqke( dq, dk, de, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T, B: tl.constexpr, @@ -587,8 +587,8 @@ def chunk_bwd_kernel_dqke( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -679,18 +679,18 @@ def chunk_ttt_linear_fwd_h( initial_state: Optional[torch.Tensor] = None, initial_state_bias: Optional[torch.Tensor] = None, output_final_state: bool = False, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BK = triton.next_power_of_2(K) BV = triton.next_power_of_2(V) assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." @@ -721,7 +721,7 @@ def chunk_ttt_linear_fwd_h( hb0=initial_state_bias, ht=final_state, hbt=final_state_bias, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, @@ -743,7 +743,7 @@ def chunk_ttt_linear_fwd_o( h: torch.Tensor, hb: torch.Tensor, scale: Optional[float] = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 64 ) -> torch.Tensor: B, T, H, K, V = *q.shape, v.shape[-1] @@ -751,8 +751,8 @@ def chunk_ttt_linear_fwd_o( scale = k.shape[-1] ** -0.5 BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = triton.next_power_of_2(K) BV = triton.next_power_of_2(V) NK = triton.cdiv(K, BK) @@ -771,8 +771,8 @@ def chunk_ttt_linear_fwd_o( h, hb, o, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -794,18 +794,18 @@ def chunk_ttt_linear_bwd_h( eps: float, initial_state: Optional[torch.Tensor] = None, initial_state_bias: Optional[torch.Tensor] = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths - if offsets is None: + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BK = triton.next_power_of_2(K) BV = triton.next_power_of_2(V) assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." @@ -836,7 +836,7 @@ def chunk_ttt_linear_bwd_h( x=x, y=y, r=rstd, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, H=H, @@ -856,14 +856,14 @@ def chunk_ttt_linear_bwd_dv_local( eta: torch.Tensor, do: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16 ) -> torch.Tensor: B, T, H, K, V = *k.shape, do.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = min(triton.next_power_of_2(K), 128) BV = min(triton.next_power_of_2(V), 128) @@ -875,8 +875,8 @@ def chunk_ttt_linear_bwd_dv_local( eta, do, dv, - offsets, - indices, + cu_seqlens, + chunk_indices, scale, T=T, H=H, @@ -908,19 +908,19 @@ def chunk_ttt_linear_bwd_norm( dv_new: Optional[torch.Tensor], # [B, H, L, D] do: torch.Tensor, # [B, H, L, D] scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # torch implementation of `dkh, dw, db, dk, dv` for LN^2 - assert offsets is None, "bwd of varlen is not implemented yet." + assert cu_seqlens is None, "bwd of varlen is not implemented yet." B, T, H, K, V = *q.shape, do.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - if offsets is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(offsets) - 1, len(indices), prepare_chunk_offsets(offsets, BT) + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) BK = triton.next_power_of_2(K) BV = triton.next_power_of_2(V) @@ -963,7 +963,7 @@ def chunk_ttt_linear_bwd_norm( dk=dk, dw=dw, db=db, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, scale=scale, T=T, @@ -996,11 +996,11 @@ def chunk_ttt_linear_bwd_norm_ref( do: torch.Tensor, # [B, H, L, D] scale: float, eps: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # torch implementation of `dkh, dw, db, dk, dv` for LN^2 - assert offsets is None, "bwd of varlen is not implemented yet." + assert cu_seqlens is None, "bwd of varlen is not implemented yet." B, T, H, K, V = *q.shape, do.shape[-1] # [B, L, H, D] -> [B, H, L, D] q, k, v, v_new, kh, y, h, eta, dv_new, do = [ @@ -1009,8 +1009,8 @@ def chunk_ttt_linear_bwd_norm_ref( ] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) pad_len = (BT - (T % BT)) % BT if pad_len > 0: q, k, v, v_new, kh, y, eta, dv_new, do = [ @@ -1092,14 +1092,14 @@ def chunk_ttt_linear_bwd_dqke( dh: torch.Tensor, dhb: torch.Tensor, scale: float, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, chunk_size: int = 16, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = chunk_size - indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = triton.next_power_of_2(K) BV = min(triton.next_power_of_2(V), 64) @@ -1123,8 +1123,8 @@ def chunk_ttt_linear_bwd_dqke( dq=dq, dk=dk, de=de, - offsets=offsets, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, scale=scale, B=B, T=T, @@ -1150,9 +1150,10 @@ def chunk_ttt_linear_fwd( initial_state: torch.Tensor, initial_state_bias: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, - BT: int = 16 + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 16 ): + BT = chunk_size h, hb, v_new, final_state, final_state_bias = chunk_ttt_linear_fwd_h( k=k, v=v, @@ -1163,7 +1164,7 @@ def chunk_ttt_linear_fwd( initial_state=initial_state, initial_state_bias=initial_state_bias, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) o = chunk_ttt_linear_fwd_o( @@ -1174,7 +1175,7 @@ def chunk_ttt_linear_fwd( h=h, hb=hb, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) return o, final_state, final_state_bias @@ -1192,11 +1193,12 @@ def chunk_ttt_linear_bwd( do: torch.Tensor, dht: torch.Tensor, dhbt: torch.Tensor, - BT: int = 16, + chunk_size: int = 16, initial_state: torch.Tensor = None, initial_state_bias: torch.Tensor = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): + BT = chunk_size h, v_new, x, y, rstd = chunk_ttt_linear_bwd_h( k=k, v=v, @@ -1206,7 +1208,7 @@ def chunk_ttt_linear_bwd( eps=eps, initial_state=initial_state, initial_state_bias=initial_state_bias, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dv_new = chunk_ttt_linear_bwd_dv_local( @@ -1215,7 +1217,7 @@ def chunk_ttt_linear_bwd( eta=eta, do=do, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dh, dhb, dh0, dhb0, dv, dk, dw, db = chunk_ttt_linear_bwd_norm( @@ -1237,7 +1239,7 @@ def chunk_ttt_linear_bwd( dv_new=dv_new, do=do, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dq, dk2, de = chunk_ttt_linear_bwd_dqke( @@ -1250,7 +1252,7 @@ def chunk_ttt_linear_bwd( dh=dh, dhb=dhb, scale=scale, - offsets=offsets, + cu_seqlens=cu_seqlens, chunk_size=BT ) dk.add_(dk2) @@ -1263,10 +1265,20 @@ class ChunkTTTLinearFunction(torch.autograd.Function): @input_guard @autocast_custom_fwd def forward( - ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state, + ctx, + q, + k, + v, + w, + b, + chunk_size, + eta, + scale, + eps, + initial_state, initial_state_bias, output_final_state, - offsets + cu_seqlens ): o, final_state, final_state_bias = chunk_ttt_linear_fwd( q=q, @@ -1277,17 +1289,17 @@ def forward( eta=eta, scale=scale, eps=eps, - BT=BT, + chunk_size=chunk_size, initial_state=initial_state, initial_state_bias=initial_state_bias, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias) - ctx.BT = BT + ctx.chunk_size = chunk_size ctx.scale = scale ctx.eps = eps - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o.to(q.dtype), final_state, final_state_bias @staticmethod @@ -1307,10 +1319,10 @@ def backward(ctx, do, dht, dhbt): do=do, dht=dht, dhbt=dhbt, - BT=ctx.BT, + chunk_size=ctx.chunk_size, initial_state=initial_state, initial_state_bias=initial_state_bias, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, ) return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None, None diff --git a/fla/ops/ttt/fused_chunk.py b/fla/ops/ttt/fused_chunk.py index d3d569c53e..3f55ba64a2 100755 --- a/fla/ops/ttt/fused_chunk.py +++ b/fla/ops/ttt/fused_chunk.py @@ -19,7 +19,7 @@ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( configs=[ @@ -44,7 +44,7 @@ def fused_chunk_ttt_linear_fwd_kernel( hb0, ht, hbt, - offsets, + cu_seqlens, T, H: tl.constexpr, K: tl.constexpr, @@ -57,11 +57,10 @@ def fused_chunk_ttt_linear_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - # indices i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 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 NT = tl.cdiv(T, BT) else: @@ -180,7 +179,6 @@ def fused_chunk_ttt_linear_bwd_kernel_h( USE_INITIAL_STATE: tl.constexpr, USE_INITIAL_STATE_B: tl.constexpr, ): - # indices i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H bos, _ = i_n * T, i_n * T + T @@ -308,7 +306,6 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( USE_FINAL_STATE_GRADIENT: tl.constexpr, USE_FINAL_STATE_GRADIENT_B: tl.constexpr, ): - # indices i_nh = tl.program_id(0) i_n, i_h = i_nh // H, i_nh % H bos, _ = i_n * T, i_n * T + T @@ -436,9 +433,9 @@ def fused_chunk_ttt_linear_bwd_h( BT: int = 16, initial_state: torch.Tensor = None, initial_state_bias: torch.Tensor = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): - assert offsets is None, "bwd of varlen is not implemented yet." + assert cu_seqlens is None, "bwd of varlen is not implemented yet." B, T, H, K, V = *k.shape, v.shape[-1] # N: the actual number of sequences in the batch with either equal or variable lengths N, NT = B, triton.cdiv(T, BT) @@ -500,9 +497,9 @@ def fused_chunk_ttt_linear_bwd_dh( BT: int = 16, initial_state: torch.Tensor = None, initial_state_bias: torch.Tensor = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): - assert offsets is None, "bwd of varlen is not implemented yet." + assert cu_seqlens is None, "bwd of varlen is not implemented yet." B, T, H, K, V = *k.shape, v.shape[-1] # N: the actual number of sequences in the batch with either equal or variable lengths N = B @@ -566,12 +563,12 @@ def fused_chunk_ttt_linear_fwd( initial_state: torch.Tensor, initial_state_bias: torch.Tensor, output_final_state: bool, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, BT: int = 16 ): B, T, H, K, V = *k.shape, v.shape[-1] # N: the actual number of sequences in the batch with either equal or variable lengths - N = B if offsets is None else len(offsets) - 1 + N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V) assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." o = torch.empty_like(v) @@ -593,7 +590,7 @@ def fused_chunk_ttt_linear_fwd( hb0=initial_state_bias, ht=final_state, hbt=final_state_bias, - offsets=offsets, + cu_seqlens=cu_seqlens, T=T, H=H, K=K, @@ -620,9 +617,9 @@ def fused_chunk_ttt_linear_bwd( BT: int = 16, initial_state: torch.Tensor = None, initial_state_bias: torch.Tensor = None, - offsets: Optional[torch.LongTensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, ): - assert offsets is None, "bwd of varlen is not implemented yet." + assert cu_seqlens is None, "bwd of varlen is not implemented yet." dq, h, v2, x, y, rstd = fused_chunk_ttt_linear_bwd_h( q=q, k=k, @@ -636,7 +633,7 @@ def fused_chunk_ttt_linear_bwd( BT=BT, initial_state=initial_state, initial_state_bias=initial_state_bias, - offsets=offsets, + cu_seqlens=cu_seqlens, ) dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd_dh( q=q, @@ -657,7 +654,7 @@ def fused_chunk_ttt_linear_bwd( BT=BT, initial_state=initial_state, initial_state_bias=initial_state_bias, - offsets=offsets, + cu_seqlens=cu_seqlens, ) return dq, dk, dv, de, dw, db, dh0, dhb0 @@ -668,7 +665,7 @@ class FusedChunkTTTLinearFunction(torch.autograd.Function): @input_guard @autocast_custom_fwd def forward(ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state, - initial_state_bias, output_final_state, offsets): + initial_state_bias, output_final_state, cu_seqlens): o, final_state, final_state_bias = fused_chunk_ttt_linear_fwd( q=q, k=k, @@ -682,13 +679,13 @@ def forward(ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state, initial_state=initial_state, initial_state_bias=initial_state_bias, output_final_state=output_final_state, - offsets=offsets, + cu_seqlens=cu_seqlens, ) ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias) ctx.BT = BT ctx.scale = scale ctx.eps = eps - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o.to(q.dtype), final_state, final_state_bias @staticmethod @@ -711,7 +708,7 @@ def backward(ctx, do, dht, dhbt): BT=ctx.BT, initial_state=initial_state, initial_state_bias=initial_state_bias, - offsets=ctx.offsets, + cu_seqlens=ctx.cu_seqlens, ) return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None diff --git a/fla/ops/utils/pooling.py b/fla/ops/utils/pooling.py index 2b1849fcd7..60f1c70003 100644 --- a/fla/ops/utils/pooling.py +++ b/fla/ops/utils/pooling.py @@ -12,7 +12,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -26,8 +26,8 @@ def mean_pooling_fwd_kernel( x, o, - offsets, - indices, + cu_seqlens, + chunk_indices, T: tl.constexpr, H: tl.constexpr, D: tl.constexpr, @@ -40,8 +40,8 @@ def mean_pooling_fwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -59,7 +59,7 @@ def mean_pooling_fwd_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -73,8 +73,8 @@ def mean_pooling_fwd_kernel( def mean_pooling_bwd_kernel( do, dx, - offsets, - indices, + cu_seqlens, + chunk_indices, T: tl.constexpr, H: tl.constexpr, D: tl.constexpr, @@ -87,8 +87,8 @@ def mean_pooling_bwd_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t - 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) + 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 NT = tl.cdiv(T, BT) else: @@ -108,20 +108,20 @@ def mean_pooling_bwd_kernel( def mean_pooling_fwd( x: torch.Tensor, chunk_size: int, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ) -> torch.Tensor: B, T, H, D = x.shape BT = chunk_size - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = x.new_empty(B, NT, H, D) def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) mean_pooling_fwd_kernel[grid]( x, o, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, D=D, @@ -136,20 +136,20 @@ def mean_pooling_bwd( batch_size: int, seq_len: int, chunk_size: int, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ) -> torch.Tensor: B, T, H, D = batch_size, seq_len, *do.shape[-2:] BT = chunk_size - indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None - NT = triton.cdiv(T, BT) if offsets is None else len(indices) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dx = do.new_empty(B, T, H, D) def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) mean_pooling_bwd_kernel[grid]( do, dx, - offsets, - indices, + cu_seqlens, + chunk_indices, T=T, H=H, D=D, @@ -168,13 +168,13 @@ def forward( ctx, x: torch.Tensor, chunk_size: int, - offsets: Optional[torch.LongTensor] = None + cu_seqlens: Optional[torch.LongTensor] = None ) -> torch.Tensor: - o = mean_pooling_fwd(x, chunk_size, offsets) + o = mean_pooling_fwd(x, chunk_size, cu_seqlens) ctx.batch_size = x.shape[0] ctx.seq_len = x.shape[1] ctx.chunk_size = chunk_size - ctx.offsets = offsets + ctx.cu_seqlens = cu_seqlens return o @staticmethod @@ -186,8 +186,8 @@ def backward( batch_size = ctx.batch_size seq_len = ctx.seq_len chunk_size = ctx.chunk_size - offsets = ctx.offsets - dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, offsets) + cu_seqlens = ctx.cu_seqlens + dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, cu_seqlens) return dx, None, None diff --git a/fla/ops/utils/solve_tril.py b/fla/ops/utils/solve_tril.py index 838044022b..f8bad9a4fe 100644 --- a/fla/ops/utils/solve_tril.py +++ b/fla/ops/utils/solve_tril.py @@ -12,7 +12,7 @@ @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -26,8 +26,8 @@ def solve_tril_16x16_kernel( A, Ad, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, BT: tl.constexpr, @@ -36,8 +36,8 @@ def solve_tril_16x16_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -62,7 +62,7 @@ def solve_tril_16x16_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -77,8 +77,8 @@ def merge_16x16_to_32x32_inverse_kernel( A, Ad, Ai, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, BT: tl.constexpr, @@ -87,8 +87,8 @@ def merge_16x16_to_32x32_inverse_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -114,7 +114,7 @@ def merge_16x16_to_32x32_inverse_kernel( @triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None }) @triton.autotune( configs=[ @@ -129,8 +129,8 @@ def merge_16x16_to_64x64_inverse_kernel( A, Ad, Ai, - offsets, - indices, + cu_seqlens, + chunk_indices, T, H: tl.constexpr, BT: tl.constexpr, @@ -139,8 +139,8 @@ def merge_16x16_to_64x64_inverse_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H 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) + 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 else: bos, eos = i_b * T, i_b * T + T @@ -246,13 +246,13 @@ def solve_tril( B, T, H, BT = A.shape Ad = torch.empty(B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype) - indices = prepare_chunk_indices(cu_seqlens, 16) if cu_seqlens is not None else None - NT = len(indices) if cu_seqlens is not None else triton.cdiv(T, 16) + chunk_indices = prepare_chunk_indices(cu_seqlens, 16) if cu_seqlens is not None else None + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 16) solve_tril_16x16_kernel[NT, B * H]( A=A, Ad=Ad, - offsets=cu_seqlens, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, BT=BT, @@ -262,14 +262,14 @@ def solve_tril( Ai = torch.zeros(B, T, H, BT, device=A.device, dtype=output_dtype) merge_fn = merge_16x16_to_32x32_inverse_kernel if BT == 32 else merge_16x16_to_64x64_inverse_kernel - indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = len(indices) if cu_seqlens is not None else triton.cdiv(T, BT) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) merge_fn[NT, B * H]( A=A, Ad=Ad, Ai=Ai, - offsets=cu_seqlens, - indices=indices, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, T=T, H=H, BT=BT, diff --git a/tests/ops/test_attn.py b/tests/ops/test_attn.py index 4ef32f4d13..219646657c 100644 --- a/tests/ops/test_attn.py +++ b/tests/ops/test_attn.py @@ -21,23 +21,23 @@ test_b_list = [2] test_t_list = [2048] test_t_varlen_list = test_t_list - test_d_list = [64, 100, 128, 256] + test_d_list = [64, 100, 128] else: test_b_list = [2, 4] test_t_list = [1, 15, 63, 286, 300, 1024, 2048] test_t_varlen_list = [63, 286, 300, 512] - test_d_list = [64, 32, 100, 256] + test_d_list = [32, 64, 100] test_hq_list = [8, 16] test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("HQ", test_hq_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.1]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('HQ', test_hq_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.1]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( not HAS_FLASH, reason="Skipping test because flash-attn is not installed" @@ -79,12 +79,12 @@ def test_parallel( assert_close("dv", ref_dv, tri_dv, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("HQ", test_hq_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('HQ', test_hq_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( not HAS_FLASH, reason="Skipping test because flash-attn is not installed" @@ -104,7 +104,7 @@ def test_parallel_varlen( N = min(1, N) if T < 64 else N # randomly split the sequence into N segments - offsets = torch.cat([ + cu_seqlens = torch.cat([ torch.tensor([0], dtype=torch.long), torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], torch.tensor([T], dtype=torch.long) @@ -119,10 +119,10 @@ def test_parallel_varlen( q=q.squeeze(0), k=k.squeeze(0), v=v.squeeze(0), - cu_seqlens_q=offsets, - cu_seqlens_k=offsets, - max_seqlen_q=prepare_lens(offsets).max(), - max_seqlen_k=prepare_lens(offsets).max(), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=prepare_lens(cu_seqlens).max(), + max_seqlen_k=prepare_lens(cu_seqlens).max(), causal=True ) ref.backward(do.squeeze(0)) @@ -134,14 +134,14 @@ def test_parallel_varlen( q=q, k=k, v=v, - cu_seqlens=offsets + cu_seqlens=cu_seqlens ) tri.backward(do) tri_dq, q.grad = q.grad.clone(), None tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005) - assert_close(" dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005) - assert_close(" dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005) + assert_close(" o", ref, tri, 0.004) + assert_close("dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005) + assert_close("dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005) + assert_close("dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005) diff --git a/tests/ops/test_based.py b/tests/ops/test_based.py index 89ec29b480..ee6b06c2d2 100644 --- a/tests/ops/test_based.py +++ b/tests/ops/test_based.py @@ -21,14 +21,14 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_based( B: int, diff --git a/tests/ops/test_delta.py b/tests/ops/test_delta.py index d35aff282e..6ad19c2dfc 100644 --- a/tests/ops/test_delta.py +++ b/tests/ops/test_delta.py @@ -23,19 +23,19 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [1]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [1]) +@pytest.mark.parametrize('dtype', [torch.bfloat16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk( B: int, @@ -80,28 +80,28 @@ def test_chunk( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad - assert_close(" o", ref, tri, 0.006) - assert_close(" ht", ref_ht, tri_ht, 0.006) - assert_close(" dq", ref_dq, tri_dq, 0.008) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.008) - assert_close(" db", ref_dbeta, tri_dbeta, 0.008) - assert_close("dh0", ref_dh0, tri_dh0, 0.008) + assert_close(' o', ref, tri, 0.006) + assert_close(' ht', ref_ht, tri_ht, 0.006) + assert_close(' dq', ref_dq, tri_dq, 0.008) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.008) + assert_close(' db', ref_dbeta, tri_dbeta, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("N", [4]) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [1]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize('N', [4]) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [1]) +@pytest.mark.parametrize('dtype', [torch.bfloat16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk_varlen( N: int, @@ -156,24 +156,24 @@ def test_chunk_varlen( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad - assert_close(" o", ref, tri, 0.005) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.008) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.008) - assert_close(" db", ref_dbeta, tri_dbeta, 0.008) - assert_close("dh0", ref_dh0, tri_dh0, 0.008) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.008) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.008) + assert_close(' db', ref_dbeta, tri_dbeta, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.1]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.1]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_l2_in_kernel( B: int, @@ -219,13 +219,13 @@ def test_l2_in_kernel( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = beta.grad = h0.grad = None - assert_close(" o", ref, tri, 0.01) - assert_close(" ht", ref_ht, tri_ht, 0.01) - assert_close(" dq", ref_dq, tri_dq, 0.01) - assert_close(" dk", ref_dk, tri_dk, 0.01) - assert_close(" dv", ref_dv, tri_dv, 0.01) - assert_close(" db", ref_dbeta, tri_dbeta, 0.01) - assert_close("dh0", ref_dh0, tri_dh0, 0.01) + assert_close(' o', ref, tri, 0.01) + assert_close(' ht', ref_ht, tri_ht, 0.01) + assert_close(' dq', ref_dq, tri_dq, 0.01) + assert_close(' dk', ref_dk, tri_dk, 0.01) + assert_close(' dv', ref_dv, tri_dv, 0.01) + assert_close(' db', ref_dbeta, tri_dbeta, 0.01) + assert_close('dh0', ref_dh0, tri_dh0, 0.01) tri, tri_ht = fused_recurrent_delta_rule( F.normalize(q.clone().float(), p=2, dim=-1).to(dtype), @@ -254,13 +254,13 @@ def test_l2_in_kernel( ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = beta.grad = h0.grad = None - assert_close(" o", ref, tri, 0.002) - assert_close(" ht", ref_ht, tri_ht, 0.002) - assert_close(" dq", ref_dq, tri_dq, 0.002) - assert_close(" dk", ref_dk, tri_dk, 0.002) - assert_close(" dv", ref_dv, tri_dv, 0.002) - assert_close(" db", ref_dbeta, tri_dbeta, 0.002) - assert_close("dh0", ref_dh0, tri_dh0, 0.002) + assert_close(' o', ref, tri, 0.002) + assert_close(' ht', ref_ht, tri_ht, 0.002) + assert_close(' dq', ref_dq, tri_dq, 0.002) + assert_close(' dk', ref_dk, tri_dk, 0.002) + assert_close(' dv', ref_dv, tri_dv, 0.002) + assert_close(' db', ref_dbeta, tri_dbeta, 0.002) + assert_close('dh0', ref_dh0, tri_dh0, 0.002) tri, tri_ht = fused_recurrent_delta_rule( F.normalize(q.float().clone(), p=2, dim=-1).to(dtype), @@ -288,10 +288,10 @@ def test_l2_in_kernel( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = beta.grad = h0.grad = None - assert_close(" o", ref, tri, 0.002) - assert_close(" ht", ref_ht, tri_ht, 0.002) - assert_close(" dq", ref_dq, tri_dq, 0.002) - assert_close(" dk", ref_dk, tri_dk, 0.002) - assert_close(" dv", ref_dv, tri_dv, 0.002) - assert_close(" db", ref_dbeta, tri_dbeta, 0.002) - assert_close("dh0", ref_dh0, tri_dh0, 0.002) + assert_close(' o', ref, tri, 0.002) + assert_close(' ht', ref_ht, tri_ht, 0.002) + assert_close(' dq', ref_dq, tri_dq, 0.002) + assert_close(' dk', ref_dk, tri_dk, 0.002) + assert_close(' dv', ref_dv, tri_dv, 0.002) + assert_close(' db', ref_dbeta, tri_dbeta, 0.002) + assert_close('dh0', ref_dh0, tri_dh0, 0.002) diff --git a/tests/ops/test_dplr_delta.py b/tests/ops/test_dplr_delta.py index f2fd712845..be839283de 100644 --- a/tests/ops/test_dplr_delta.py +++ b/tests/ops/test_dplr_delta.py @@ -36,28 +36,21 @@ def recurrent_dplr_delta_rule_ref( scale: float = None, initial_state: torch.Tensor = None, output_final_state: bool = False, - head_first=False, ): - q, k, v, a, b, gk = map(lambda x: x.to(torch.float32), [q, k, v, a, b, gk]) - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - a = a.transpose(1, 2) - b = b.transpose(1, 2) - gk = gk.transpose(1, 2) - B, H, L, DK = q.shape - DV = v.shape[-1] + q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk)) + + B, H, T, K, V = *q.shape, v.shape[-1] o = torch.zeros_like(v) - S = torch.zeros(B, H, DK, DV).to(v) + S = torch.zeros(B, H, K, V).to(v) if initial_state is not None: S = initial_state if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) + scale = K ** -0.5 q = q * scale - for i in range(L): - _k = k[:, :, i] + + for i in range(T): _q = q[:, :, i] + _k = k[:, :, i] _v = v[:, :, i].clone() a_i = a[:, :, i] b_i = b[:, :, i] @@ -68,8 +61,7 @@ def recurrent_dplr_delta_rule_ref( o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) if not output_final_state: S = None - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o, S @@ -84,59 +76,57 @@ def chunk_dplr_delta_rule_ref( output_final_state: bool = True, scale: float = None, chunk_size: int = 64, - head_first: bool = True, ): + q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk)) BT = chunk_size - if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) - if not head_first: - q, k, v, a, b, gk = map(lambda x: rearrange(x, 'b t h ... -> b h t ...'), (q, k, v, a, b, gk)) T = q.shape[-2] pad_len = (BT - (T % BT)) % BT - q, k, v, a, b, gk = map(lambda x: F.pad(x, (0, 0, 0, pad_len)).to(torch.float32), [q, k, v, a, b, gk]) - B, H, L, DK = q.shape - DV = v.shape[-1] + q, k, v, a, b, gk = map(lambda x: F.pad(x, (0, 0, 0, pad_len)).to(torch.float), [q, k, v, a, b, gk]) + B, H, _, K, V = *q.shape, v.shape[-1] + NT = q.shape[-2] // BT + if scale is None: + scale = K ** -0.5 q = q * scale - S = k.new_zeros(B, H, DK, DV) + S = k.new_zeros(B, H, K, V) if initial_state is not None: S += initial_state # note that diagonal is masked. - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, a, b, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, a, b, gk]) + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, a, b, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, v, a, b, gk]) gk_cumsum = gk.cumsum(-2) - A_ab = torch.zeros(B, H, L // chunk_size, chunk_size, chunk_size).to(q.device) - A_qk = torch.zeros(B, H, L // chunk_size, chunk_size, chunk_size).to(q.device) - A_ak = torch.zeros(B, H, L // chunk_size, chunk_size, chunk_size).to(q.device) - A_qb = torch.zeros(B, H, L // chunk_size, chunk_size, chunk_size).to(q.device) + A_ab = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_qk = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_ak = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_qb = torch.zeros(B, H, NT, BT, BT).to(q.device) - for i in range(chunk_size): + for i in range(BT): a_i = a[:, :, :, i, None] q_i = q[:, :, :, i, None] gk_i = gk_cumsum[:, :, :, i, None] - mask = (torch.arange(chunk_size) <= i).to(q.device) + mask = (torch.arange(BT) <= i).to(q.device) attn_i = (gk_i - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() A_qk[:, :, :, i, :] = (q_i * k * attn_i).sum(-1).clone() A_qb[:, :, :, i, :] = (q_i * b * attn_i).sum(-1).clone() - mask = (torch.arange(chunk_size) < i).to(q.device) + mask = (torch.arange(BT) < i).to(q.device) # shift by one. attn_i = (gk_i - gk[:, :, :, i, None] - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() A_ab[:, :, :, i, :] = (a_i * b * attn_i).sum(-1).clone() A_ak[:, :, :, i, :] = (a_i * k * attn_i).sum(-1).clone() A_ab = A_ab - for i in range(1, chunk_size): + for i in range(1, BT): A_ab[..., i, :i] = A_ab[..., i, :i].clone() + (A_ab[..., i, :, None].clone() * A_ab[..., :, :i].clone()).sum(-2) - A_ab = A_ab + torch.eye(chunk_size, dtype=torch.float, device=q.device) + A_ab = A_ab + torch.eye(BT, dtype=torch.float, device=q.device) u = A_ab @ (A_ak @ v) w = A_ab @ ((gk_cumsum-gk).exp() * a) o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, L // chunk_size): + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, NT): q_i, k_i, v_i, u_i, w_i, b_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], b[:, :, i] v2_i = u_i + w_i @ S o_1 = A_qk[:, :, i] @ v_i @@ -149,23 +139,21 @@ def chunk_dplr_delta_rule_ref( S = None if output_final_state is False else S o = rearrange(o, 'b h n c d -> b h (n c) d') - o = o[:, :, :T] - if not head_first: - o = o.transpose(1, 2) + o = o[:, :, :T].transpose(1, 2) return o, S -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) -def test_recurrent_forward( +def test_recurrent_fwd( B: int, T: int, H: int, @@ -173,27 +161,20 @@ def test_recurrent_forward( scale: float, dtype: torch.dtype, ): - head_first = True + torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' os.environ['TORCH_CUDA_MATMUL_PRECISION'] = 'highest' - if head_first: - q = torch.randn(B, H, T, D, dtype=dtype) - k = torch.randn(B, H, T, D, dtype=dtype) - v = torch.randn(B, H, T, D, dtype=dtype) - a = torch.rand(B, H, T, D, dtype=dtype) - gk = (torch.randn(B, H, T, D, dtype=torch.float)) - else: - q = torch.randn(B, T, H, D, dtype=dtype) - k = torch.randn(B, T, H, D, dtype=dtype) - v = torch.randn(B, T, H, D, dtype=dtype) - a = torch.rand(B, T, H, D, dtype=dtype) - gk = torch.randn(B, T, H, D, dtype=torch.float) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) a = F.normalize(a, p=2, dim=-1) b = -a - gk = torch.nn.functional.logsigmoid(gk) / 16 + gk = F.logsigmoid(gk) / 16 - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + h0 = torch.randn(B, H, D, D, dtype=torch.float) q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0)) ref, ref_ht = chunk_dplr_delta_rule_ref( q=q.clone(), @@ -205,7 +186,6 @@ def test_recurrent_forward( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) tri, tri_ht = recurrent_dplr_delta_rule_ref( q=q.clone(), @@ -217,23 +197,20 @@ def test_recurrent_forward( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) - assert_close(" o", ref, tri, 0.001) - assert_close(" ht", ref_ht, tri_ht, 0.001) - - -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) -@pytest.mark.parametrize("head_first", [True, False]) -@pytest.mark.parametrize("compile", [False, True]) + assert_close(' o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + + +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_fused_recurrent_fwd( B: int, @@ -242,27 +219,19 @@ def test_fused_recurrent_fwd( D: int, scale: float, dtype: torch.dtype, - head_first: bool, - compile: bool, ): - if head_first: - q = torch.randn(B, H, T, D, dtype=dtype) - k = torch.randn(B, H, T, D, dtype=dtype) - v = torch.randn(B, H, T, D, dtype=dtype) - a = torch.rand(B, H, T, D, dtype=dtype) - gk = (torch.randn(B, H, T, D, dtype=torch.float)) - else: - q = torch.randn(B, T, H, D, dtype=dtype) - k = torch.randn(B, T, H, D, dtype=dtype) - v = torch.randn(B, T, H, D, dtype=dtype) - a = torch.rand(B, T, H, D, dtype=dtype) - gk = torch.randn(B, T, H, D, dtype=torch.float) + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) a = F.normalize(a, p=2, dim=-1) b = -a - gk = torch.nn.functional.logsigmoid(gk) / 4 + gk = F.logsigmoid(gk) / 4 - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + h0 = torch.randn(B, H, D, D, dtype=torch.float) q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0)) ref, ref_ht = recurrent_dplr_delta_rule_ref( q=q.clone(), @@ -274,12 +243,9 @@ def test_fused_recurrent_fwd( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) - fused_compiled = torch.compile(fused_recurrent_dplr_delta_rule) if compile else fused_recurrent_dplr_delta_rule - - tri, tri_ht = fused_compiled( + tri, tri_ht = fused_recurrent_dplr_delta_rule( q=q.clone(), k=k.clone(), v=v.clone(), @@ -289,28 +255,26 @@ def test_fused_recurrent_fwd( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) - assert_close(" o", ref, tri, 0.002) - assert_close(" ht", ref_ht, tri_ht, 0.002) - - -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) -@pytest.mark.parametrize("head_first", [False, True]) -@pytest.mark.parametrize("compile", [False, True]) + assert_close(' o', ref, tri, 0.002) + assert_close('ht', ref_ht, tri_ht, 0.002) + + +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) +@pytest.mark.parametrize('compile', [False, True]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk( B: int, @@ -320,27 +284,20 @@ def test_chunk( scale: float, gate_logit_normalizer: float, dtype: torch.dtype, - head_first: bool, compile: bool, ): - if head_first: - q = torch.randn(B, H, T, D, dtype=dtype) - k = torch.randn(B, H, T, D, dtype=dtype) - v = torch.randn(B, H, T, D, dtype=dtype) - a = torch.rand(B, H, T, D, dtype=dtype) - gk = (torch.randn(B, H, T, D, dtype=torch.float)) - else: - q = torch.randn(B, T, H, D, dtype=dtype) - k = torch.randn(B, T, H, D, dtype=dtype) - v = torch.randn(B, T, H, D, dtype=dtype) - a = torch.rand(B, T, H, D, dtype=dtype) - gk = torch.randn(B, T, H, D, dtype=torch.float) + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) a = F.normalize(a, p=2, dim=-1) b = -a gk = F.logsigmoid(gk) / gate_logit_normalizer - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + h0 = torch.randn(B, H, D, D, dtype=torch.float) q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0)) ref, ref_ht = chunk_dplr_delta_rule_ref( q=q.clone(), @@ -352,7 +309,6 @@ def test_chunk( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) do = torch.randn_like(v) dht = torch.randn_like(h0) @@ -372,37 +328,36 @@ def test_chunk( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) tri_dq, tri_dk, tri_dv, tri_da, tri_db, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None - assert_close(" o", ref, tri, 0.007) - assert_close(" ht", ref_ht, tri_ht, 0.008) - assert_close(" dq", ref_dq, tri_dq, 0.008) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.008) - assert_close(" da", ref_da, tri_da, 0.008) - assert_close(" db", ref_db, tri_db, 0.008) + assert_close(' o', ref, tri, 0.007) + assert_close(' ht', ref_ht, tri_ht, 0.008) + assert_close(' dq', ref_dq, tri_dq, 0.008) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.008) + assert_close(' da', ref_da, tri_da, 0.008) + assert_close(' db', ref_db, tri_db, 0.008) if gate_logit_normalizer >= 1 and ref_dg.norm() > 0.01: # otherwise it is meaningless - assert_close(" dg", ref_dg, tri_dg, 0.008) - assert_close("dh0", ref_dh0, tri_dh0, 0.008) + assert_close(' dg', ref_dg, tri_dg, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk_varlen( N: int, @@ -415,7 +370,7 @@ def test_chunk_varlen( torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments - offsets = torch.cat([ + cu_seqlens = torch.cat([ torch.tensor([0], dtype=torch.long), torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], torch.tensor([T], dtype=torch.long) @@ -428,8 +383,8 @@ def test_chunk_varlen( gk = torch.randn(1, T, H, D, dtype=torch.float) a = F.normalize(a, p=2, dim=-1) b = -a - gk = torch.nn.functional.logsigmoid(gk) - h0 = torch.randn(N, H, D, D, dtype=torch.float32) + gk = F.logsigmoid(gk) + h0 = torch.randn(N, H, D, D, dtype=torch.float) q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0)) tri, tri_ht = chunk_dplr_delta_rule( @@ -442,8 +397,7 @@ def test_chunk_varlen( scale=scale, output_final_state=True, initial_state=h0.clone(), - cu_seqlens=offsets, - head_first=False + cu_seqlens=cu_seqlens, ) do = torch.randn_like(v) dht = torch.randn_like(h0) @@ -455,16 +409,15 @@ def test_chunk_varlen( ref_ht = [] for i in range(N): ref_i, ref_ht_i = chunk_dplr_delta_rule_ref( - q=q[:, offsets[i]:offsets[i+1]], - k=k[:, offsets[i]:offsets[i+1]], - v=v[:, offsets[i]:offsets[i+1]], - a=a[:, offsets[i]:offsets[i+1]], - b=b[:, offsets[i]:offsets[i+1]], - gk=gk[:, offsets[i]:offsets[i+1]], + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + a=a[:, cu_seqlens[i]:cu_seqlens[i+1]], + b=b[:, cu_seqlens[i]:cu_seqlens[i+1]], + gk=gk[:, cu_seqlens[i]:cu_seqlens[i+1]], scale=scale, initial_state=h0[i, None], output_final_state=True, - head_first=False ) ref.append(ref_i) ref_ht.append(ref_ht_i) @@ -474,12 +427,12 @@ def test_chunk_varlen( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_da, ref_db, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad - assert_close(" o", ref, tri, 0.007) - assert_close(" ht", ref_ht, tri_ht, 0.008) - assert_close(" dq", ref_dq, tri_dq, 0.008) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.008) - assert_close(" da", ref_da, tri_da, 0.008) - assert_close(" db", ref_db, tri_db, 0.008) - assert_close(" dg", ref_dg, tri_dg, 0.008) - assert_close("dh0", ref_dh0, tri_dh0, 0.008) + assert_close(' o', ref, tri, 0.007) + assert_close(' ht', ref_ht, tri_ht, 0.008) + assert_close(' dq', ref_dq, tri_dq, 0.008) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.008) + assert_close(' da', ref_da, tri_da, 0.008) + assert_close(' db', ref_db, tri_db, 0.008) + assert_close(' dg', ref_dg, tri_dg, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) diff --git a/tests/ops/test_gated_delta.py b/tests/ops/test_gated_delta.py index f8a7d295bd..b79317fc7f 100644 --- a/tests/ops/test_gated_delta.py +++ b/tests/ops/test_gated_delta.py @@ -141,17 +141,17 @@ def chunk_gated_delta_rule_ref( return o, S -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("scale", [1]) -@pytest.mark.parametrize("use_qk_l2norm_in_kernel", [True, False]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('scale', [1]) +@pytest.mark.parametrize('use_qk_l2norm_in_kernel', [True, False]) +@pytest.mark.parametrize('dtype', [torch.float32, torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_recurrent_forward( B: int, @@ -193,20 +193,20 @@ def test_recurrent_forward( use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, output_final_state=True, ) - assert_close(" o", ref, tri, 0.002) - assert_close(" ht", ref_ht, tri_ht, 0.002) + assert_close(' o', ref, tri, 0.002) + assert_close(' ht', ref_ht, tri_ht, 0.002) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("scale", [1, 0.1]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('scale', [1, 0.1]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -218,7 +218,7 @@ def test_chunk( gate_logit_normalizer: float ): if is_intel_alchemist and D > 128: - pytest.skip(reason="chunk_gated_delta_rule is not supported on alchemist for D>128") + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') q = torch.randn(B, T, H, D, dtype=dtype) k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) @@ -258,26 +258,26 @@ def test_chunk( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad - assert_close(" o", ref, tri, 0.005) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.008) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.008) - assert_close(" db", ref_dbeta, tri_dbeta, 0.02) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.008) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.008) + assert_close(' db', ref_dbeta, tri_dbeta, 0.02) if gate_logit_normalizer >= 1 and ref_dg.norm() > 0.01: - assert_close(" dg", ref_dg, tri_dg, 0.02) - assert_close("dh0", ref_dh0, tri_dh0, 0.008) + assert_close(' dg', ref_dg, tri_dg, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [1, 0.1]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [1, 0.1]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_chunk_varlen( N: int, @@ -288,7 +288,7 @@ def test_chunk_varlen( dtype: torch.dtype, ): if is_intel_alchemist and D > 128: - pytest.skip(reason="chunk_gated_delta_rule is not supported on alchemist for D>128") + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments @@ -345,11 +345,11 @@ def test_chunk_varlen( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad - assert_close(" o", ref, tri, 0.005) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.007) - assert_close(" dk", ref_dk, tri_dk, 0.008) - assert_close(" dv", ref_dv, tri_dv, 0.007) - assert_close(" db", ref_dbeta, tri_dbeta, 0.015) - assert_close(" dg", ref_dg, tri_dg, 0.015) - assert_close("dh0", ref_dh0, tri_dh0, 0.007) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.007) + assert_close(' dk', ref_dk, tri_dk, 0.008) + assert_close(' dv', ref_dv, tri_dv, 0.007) + assert_close(' db', ref_dbeta, tri_dbeta, 0.015) + assert_close(' dg', ref_dg, tri_dg, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) diff --git a/tests/ops/test_gla.py b/tests/ops/test_gla.py index 4bd671e238..93a44cdf93 100644 --- a/tests/ops/test_gla.py +++ b/tests/ops/test_gla.py @@ -26,18 +26,18 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_fused_recurrent( B: int, @@ -85,28 +85,28 @@ def test_fused_recurrent( tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) - - -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.bfloat16]) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk( B: int, @@ -165,23 +165,23 @@ def test_chunk( ref_dg, g.grad = g.grad.clone(), None ref_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_chunk_varlen( N: int, @@ -242,10 +242,10 @@ def test_chunk_varlen( tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) diff --git a/tests/ops/test_gsa.py b/tests/ops/test_gsa.py index 6eca47e88b..f81ff2de19 100644 --- a/tests/ops/test_gsa.py +++ b/tests/ops/test_gsa.py @@ -28,19 +28,19 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("M", test_m_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('M', test_m_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_fused_recurrent( B: int, @@ -98,31 +98,31 @@ def test_fused_recurrent( tri_dhk0, hk0.grad = hk0.grad.clone(), None tri_dhv0, hv0.grad = hv0.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close(" hkt", ref_hkt, tri_hkt, 0.005) - assert_close(" hvt", ref_hvt, tri_hvt, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" ds", ref_ds, tri_ds, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dhk0", ref_dhk0, tri_dhk0, 0.005) - assert_close("dhv0", ref_dhv0, tri_dhv0, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close(' hkt', ref_hkt, tri_hkt, 0.005) + assert_close(' hvt', ref_hvt, tri_hvt, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' ds', ref_ds, tri_ds, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("M", test_m_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('M', test_m_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_fused_recurrent_varlen( N: int, @@ -205,32 +205,32 @@ def test_fused_recurrent_varlen( tri_dhk0, hk0.grad = hk0.grad.clone(), None tri_dhv0, hv0.grad = hv0.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close(" hkt", ref_hkt, tri_hkt, 0.005) - assert_close(" hvt", ref_hvt, tri_hvt, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" ds", ref_ds, tri_ds, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dhk0", ref_dhk0, tri_dhk0, 0.005) - assert_close("dhv0", ref_dhv0, tri_dhv0, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close(' hkt', ref_hkt, tri_hkt, 0.005) + assert_close(' hvt', ref_hvt, tri_hvt, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' ds', ref_ds, tri_ds, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("M", test_m_list) -@pytest.mark.parametrize("dtype", [torch.float]) -@pytest.mark.parametrize("gate_logit_normalizer", [1, 0.05, 20]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('M', test_m_list) +@pytest.mark.parametrize('dtype', [torch.float]) +@pytest.mark.parametrize('gate_logit_normalizer', [1, 0.05, 20]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk( B: int, @@ -242,7 +242,7 @@ def test_chunk( gate_logit_normalizer: float, ): if (D > 64 or M > 64) and check_shared_mem('hopper') is False: - pytest.skip(reason="Current CI do not support this config") + pytest.skip(reason='Current CI do not support this config') torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' @@ -271,27 +271,27 @@ def test_chunk( tri_ds, s.grad = s.grad.clone(), None tri_dg, s.grad = g.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close("dq", ref_dq, tri_dq, 0.005) - assert_close("dk", ref_dk, tri_dk, 0.005) - assert_close("dv", ref_dv, tri_dv, 0.005) - assert_close("ds", ref_ds, tri_ds, 0.008) - assert_close("dg", ref_dg, tri_dg, 0.008) + assert_close(' o', ref, tri, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('ds', ref_ds, tri_ds, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.008) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("M", test_m_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('M', test_m_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_chunk_varlen( N: int, @@ -302,7 +302,7 @@ def test_chunk_varlen( dtype: torch.dtype, ): if (D > 64 or M > 64) and check_shared_mem('hopper') is False: - pytest.skip(reason="Current CI do not support this config") + pytest.skip(reason='Current CI do not support this config') torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments @@ -357,32 +357,32 @@ def test_chunk_varlen( tri_dhk0, hk0.grad = hk0.grad.clone(), None tri_dhv0, hv0.grad = hv0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" hkt", ref_hkt, tri_hkt, 0.005) - assert_close(" hvt", ref_hvt, tri_hvt, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" ds", ref_ds, tri_ds, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dhk0", ref_dhk0, tri_dhk0, 0.005) - assert_close("dhv0", ref_dhv0, tri_dhv0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' hkt', ref_hkt, tri_hkt, 0.005) + assert_close(' hvt', ref_hvt, tri_hvt, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' ds', ref_ds, tri_ds, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) -@pytest.mark.parametrize("HQ", [8, 16]) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("M", test_m_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('HQ', [8, 16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('M', test_m_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Triton Failure" + reason='Intel Triton Failure' ) def test_inference( B: int, @@ -416,5 +416,5 @@ def test_inference( output_final_state=True ) tri[:, i] = o.squeeze(1) - assert_close(f"o{i}", ref[:, i], tri[:, i], 0.005) + assert_close(f'o{i}', ref[:, i], tri[:, i], 0.005) h0 = ht diff --git a/tests/ops/test_hgrn.py b/tests/ops/test_hgrn.py index 0c0a9ec6b6..72c82f82e4 100644 --- a/tests/ops/test_hgrn.py +++ b/tests/ops/test_hgrn.py @@ -26,13 +26,13 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_fused_recurrent( B: int, @@ -63,20 +63,20 @@ def test_fused_recurrent( tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close("o", ref, tri, 0.005) - assert_close("ht", ref_ht, tri_ht, 0.005) - assert_close("dx", ref_dx, tri_dx, 0.005) - assert_close("dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dx', ref_dx, tri_dx, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_fused_recurrent_varlen( N: int, @@ -87,9 +87,9 @@ def test_fused_recurrent_varlen( torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments - offsets = torch.cat([ + cu_seqlens = torch.cat([ torch.tensor([0], dtype=torch.long), - torch.arange(16, T)[torch.randperm(T - 1)[:N-1]], + torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], torch.tensor([T], dtype=torch.long) ], 0).to(device).sort()[0] @@ -104,8 +104,8 @@ def test_fused_recurrent_varlen( refs, ref_hts = [], [] for i in range(N): ref, ref_ht = naive_recurrent_hgrn( - x[:, offsets[i]:offsets[i+1]], - g[:, offsets[i]:offsets[i+1]], + x[:, cu_seqlens[i]:cu_seqlens[i+1]], + g[:, cu_seqlens[i]:cu_seqlens[i+1]], h0[i:i+1], output_final_state=True ) @@ -118,26 +118,26 @@ def test_fused_recurrent_varlen( ref_dg, g.grad = g.grad.clone(), None ref_dh0, h0.grad = h0.grad.clone(), None - tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True, cu_seqlens=offsets) + tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True, cu_seqlens=cu_seqlens) ((tri * do).sum() + (tri_ht * dht).sum()).backward() tri_dx, x.grad = x.grad.clone(), None tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close("ht", ref_ht, tri_ht, 0.005) - assert_close("dx", ref_dx, tri_dx, 0.005) - assert_close("dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dx', ref_dx, tri_dx, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -165,6 +165,6 @@ def test_chunk( tri_dx, x.grad = x.grad.clone(), None tri_dg, g.grad = g.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close("dx", ref_dx, tri_dx, 0.005) - assert_close("dg", ref_dg, tri_dg, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close('dx', ref_dx, tri_dx, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) diff --git a/tests/ops/test_iplr_delta.py b/tests/ops/test_iplr_delta.py index 1bd052a40f..2ba982a89b 100644 --- a/tests/ops/test_iplr_delta.py +++ b/tests/ops/test_iplr_delta.py @@ -38,17 +38,12 @@ def chunk_iplr_delta_rule_ref( output_final_state: bool = True, scale: float = None, chunk_size: int = 64, - head_first: bool = True, ): BT = chunk_size if scale is None: scale = 1 / (q.shape[-1] ** 0.5) - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - a = a.transpose(1, 2) - b = b.transpose(1, 2) + + q, k, v, a, b = map(lambda x: x.transpose(1, 2), (q, k, v, a, b)) T = q.shape[-2] pad_len = (BT - (T % BT)) % BT if pad_len > 0: @@ -97,8 +92,7 @@ def chunk_iplr_delta_rule_ref( S = None if output_final_state is False else S o = rearrange(o, 'b h n c d -> b h (n c) d') o = o[:, :, :T] - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o, S @@ -110,19 +104,12 @@ def recurrence_iplr_delta_rule_ref( b, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = True, - head_first: bool = True, scale: Optional[float] = None ): orig_dtype = q.dtype if scale is None: scale = 1 / (q.shape[-1] ** 0.5) - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - a = a.transpose(1, 2) - b = b.transpose(1, 2) - q, k, v, a, b = map(lambda x: x.to(torch.float32), [q, k, v, a, b]) + q, k, v, a, b = map(lambda x: x.transpose(1, 2).to(torch.float32), [q, k, v, a, b]) q = q * scale B, H, L, DK = q.shape DV = v.shape[-1] @@ -141,21 +128,19 @@ def recurrence_iplr_delta_rule_ref( S = S + _kv o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) S = None if output_final_state is False else S - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o.to(orig_dtype), S -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) -@pytest.mark.parametrize("head_first", [True, False]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -164,18 +149,11 @@ def test_chunk( D: int, scale: float, dtype: torch.dtype, - head_first: bool, ): - if head_first: - q = torch.randn(B, H, T, D, dtype=dtype) - k = torch.randn(B, H, T, D, dtype=dtype) - v = torch.randn(B, H, T, D, dtype=dtype) - a = torch.rand(B, H, T, D, dtype=dtype) - else: - q = torch.randn(B, T, H, D, dtype=dtype) - k = torch.randn(B, T, H, D, dtype=dtype) - v = torch.randn(B, T, H, D, dtype=dtype) - a = torch.rand(B, T, H, D, dtype=dtype) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) a = F.normalize(a, p=2, dim=-1) b = -a @@ -190,7 +168,6 @@ def test_chunk( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) tri, tri_ht = chunk_iplr_delta_rule( q=q.clone(), @@ -201,19 +178,17 @@ def test_chunk( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) - assert_close(" o", ref, tri, 0.007) - assert_close(" ht", ref_ht, tri_ht, 0.008) + assert_close(' o', ref, tri, 0.007) + assert_close('ht', ref_ht, tri_ht, 0.008) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("scale", [0.25]) -@pytest.mark.parametrize("dtype", [torch.float16]) -@pytest.mark.parametrize("head_first", [True, False]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('scale', [0.25]) +@pytest.mark.parametrize('dtype', [torch.float16]) def test_recurrent( B: int, T: int, @@ -221,18 +196,11 @@ def test_recurrent( D: int, scale: float, dtype: torch.dtype, - head_first: bool, ): - if head_first: - q = torch.randn(B, H, T, D, dtype=dtype) - k = torch.randn(B, H, T, D, dtype=dtype) - v = torch.randn(B, H, T, D, dtype=dtype) - a = torch.rand(B, H, T, D, dtype=dtype) - else: - q = torch.randn(B, T, H, D, dtype=dtype) - k = torch.randn(B, T, H, D, dtype=dtype) - v = torch.randn(B, T, H, D, dtype=dtype) - a = torch.rand(B, T, H, D, dtype=dtype) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) a = F.normalize(a, p=2, dim=-1) b = -a @@ -247,7 +215,6 @@ def test_recurrent( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) dht = torch.rand_like(h0) do = torch.rand_like(ref) @@ -263,14 +230,13 @@ def test_recurrent( scale=scale, initial_state=h0.clone(), output_final_state=True, - head_first=head_first ) ((dht * tri_ht).sum() + (do * tri).sum()).backward() - assert_close(" o", ref, tri, 0.003) - assert_close(" ht", ref_ht, tri_ht, 0.003) - assert_close(" dq", dq, q.grad, 0.003) - assert_close(" dk", dk, k.grad, 0.003) - assert_close(" dv", dv, v.grad, 0.003) - assert_close(" da", da, a.grad, 0.003) - assert_close(" db", db, b.grad, 0.003) - assert_close("dh0", dh0, h0.grad, 0.003) + assert_close(' o', ref, tri, 0.003) + assert_close(' ht', ref_ht, tri_ht, 0.003) + assert_close(' dq', dq, q.grad, 0.003) + assert_close(' dk', dk, k.grad, 0.003) + assert_close(' dv', dv, v.grad, 0.003) + assert_close(' da', da, a.grad, 0.003) + assert_close(' db', db, b.grad, 0.003) + assert_close('dh0', dh0, h0.grad, 0.003) diff --git a/tests/ops/test_linear_attn.py b/tests/ops/test_linear_attn.py index 99326cb0ff..61c483dba5 100644 --- a/tests/ops/test_linear_attn.py +++ b/tests/ops/test_linear_attn.py @@ -25,14 +25,14 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_fused_recurrent( B: int, @@ -59,20 +59,20 @@ def test_fused_recurrent( tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.001) - assert_close(" dq", ref_dq, tri_dq, 0.001) - assert_close(" dk", ref_dk, tri_dk, 0.001) - assert_close(" dv", ref_dv, tri_dv, 0.001) + assert_close(' o', ref, tri, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -115,21 +115,21 @@ def test_chunk( tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.001) - assert_close(" ht", ref_ht, tri_ht, 0.001) - assert_close(" dq", ref_dq, tri_dq, 0.001) - assert_close(" dk", ref_dk, tri_dk, 0.001) - assert_close(" dv", ref_dv, tri_dv, 0.001) + assert_close(' o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_fused_chunk( B: int, @@ -172,8 +172,8 @@ def test_fused_chunk( tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.001) - assert_close(" ht", ref_ht, tri_ht, 0.001) - assert_close(" dq", ref_dq, tri_dq, 0.001) - assert_close(" dk", ref_dk, tri_dk, 0.001) - assert_close(" dv", ref_dv, tri_dv, 0.001) + assert_close(' o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) diff --git a/tests/ops/test_nsa.py b/tests/ops/test_nsa.py index 15dde1bf04..1f37c748ab 100644 --- a/tests/ops/test_nsa.py +++ b/tests/ops/test_nsa.py @@ -24,22 +24,22 @@ # FIXME -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("HQ", [64]) -@pytest.mark.parametrize("D", [100, 64]) -@pytest.mark.parametrize("S", [16]) -@pytest.mark.parametrize("block_size", [32]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("scale", [0.1]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('HQ', [64]) +@pytest.mark.parametrize('D', [100, 64]) +@pytest.mark.parametrize('S', [16]) +@pytest.mark.parametrize('block_size', [32]) +@pytest.mark.parametrize('dtype', [torch.float16]) +@pytest.mark.parametrize('scale', [0.1]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( True, - reason="TBD" + reason='TBD' ) def test_parallel( B: int, @@ -60,21 +60,21 @@ def test_parallel( v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) do = torch.randn((B, T, HQ, D), dtype=dtype, device=device) - indices = torch.full((B, T, H, S), T, dtype=torch.long, device=device) + block_indices = torch.full((B, T, H, S), T, dtype=torch.long, device=device) for b in range(B): for t in range(T): for h in range(H): i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S] - indices[b, t, h, :len(i_i)] = i_i - indices = indices.sort(-1)[0] + block_indices[b, t, h, :len(i_i)] = i_i + block_indices = block_indices.sort(-1)[0] - ref = naive_nsa(q=q, k=k, v=v, indices=indices, block_size=block_size, scale=scale) + ref = naive_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale) ref.backward(do) ref_dq, q.grad = q.grad.clone(), None ref_dk, k.grad = k.grad.clone(), None ref_dv, v.grad = v.grad.clone(), None - tri = parallel_nsa(q=q, k=k, v=v, indices=indices, block_size=block_size, scale=scale) + tri = parallel_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale) tri.backward(do) tri_dq, q.grad = q.grad.clone(), None tri_dk, k.grad = k.grad.clone(), None @@ -86,21 +86,21 @@ def test_parallel( assert_close("dv", ref_dv, tri_dv, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("HQ", [64]) -@pytest.mark.parametrize("D", [100, 64]) -@pytest.mark.parametrize("S", [16]) -@pytest.mark.parametrize("block_size", [32]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('HQ', [64]) +@pytest.mark.parametrize('D', [100, 64]) +@pytest.mark.parametrize('S', [16]) +@pytest.mark.parametrize('block_size', [32]) +@pytest.mark.parametrize('dtype', [torch.bfloat16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test because SKIP_TEST_CHUNK_VARLEN is set' ) @pytest.mark.skipif( True, - reason="TBD" + reason='TBD' ) def test_parallel_varlen( N: int, @@ -116,9 +116,9 @@ def test_parallel_varlen( os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments - offsets = torch.cat([ + cu_seqlens = torch.cat([ torch.tensor([0], dtype=torch.long), - torch.arange(16, T)[torch.randperm(T - 1)[:N-1]], + torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], torch.tensor([T], dtype=torch.long) ], 0).to(device).sort()[0] # seq-first required for inputs with variable lengths @@ -127,23 +127,23 @@ def test_parallel_varlen( v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() do = torch.randn((1, T, HQ, D), dtype=dtype, device=device) - indices = torch.full((1, T, H, S), T, dtype=torch.long, device=device) - seq_indices = prepare_token_indices(offsets).tolist() + block_indices = torch.full((1, T, H, S), T, dtype=torch.long, device=device) + seq_indices = prepare_token_indices(cu_seqlens).tolist() for i in range(T): _, t = seq_indices[i] for h in range(H): i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S] - indices[0, i, h, :len(i_i)] = i_i - indices = indices.sort(-1)[0] + block_indices[0, i, h, :len(i_i)] = i_i + block_indices = block_indices.sort(-1)[0] ref = naive_nsa( q=q, k=k, v=v, - indices=indices, + block_indices=block_indices, block_size=block_size, - cu_seqlens=offsets + cu_seqlens=cu_seqlens ) ref.backward(do) ref_dq, q.grad = q.grad.clone(), None @@ -154,16 +154,16 @@ def test_parallel_varlen( q=q, k=k, v=v, - indices=indices, + block_indices=block_indices, block_size=block_size, - cu_seqlens=offsets + cu_seqlens=cu_seqlens ) tri.backward(do) tri_dq, q.grad = q.grad.clone(), None tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close("dq", ref_dq, tri_dq, 0.005) - assert_close("dk", ref_dk, tri_dk, 0.005) - assert_close("dv", ref_dv, tri_dv, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) diff --git a/tests/ops/test_retention.py b/tests/ops/test_retention.py index fc1462a9b2..4fe427f7ae 100644 --- a/tests/ops/test_retention.py +++ b/tests/ops/test_retention.py @@ -23,15 +23,15 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("K", test_d_list) -@pytest.mark.parametrize("expand_ratio", [1, 2]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('K', test_d_list) +@pytest.mark.parametrize('expand_ratio', [1, 2]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -64,22 +64,22 @@ def test_chunk( tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close("ht", ref_ht, tri_ht, 0.005) - assert_close("dq", ref_dq, tri_dq, 0.005) - assert_close("dk", ref_dk, tri_dk, 0.005) - assert_close("dv", ref_dv, tri_dv, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("K", test_d_list) -@pytest.mark.parametrize("expand_ratio", [1, 2]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('K', test_d_list) +@pytest.mark.parametrize('expand_ratio', [1, 2]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_chunk_varlen( N: int, @@ -135,23 +135,23 @@ def test_chunk_varlen( tri_dv, v.grad = v.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("K", test_d_list) -@pytest.mark.parametrize("expand_ratio", [1, 2]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('K', test_d_list) +@pytest.mark.parametrize('expand_ratio', [1, 2]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_parallel( B: int, @@ -181,7 +181,7 @@ def test_parallel( tri_dk, k.grad = k.grad.clone(), None tri_dv, v.grad = v.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close("dq", ref_dq, tri_dq, 0.005) - assert_close("dk", ref_dk, tri_dk, 0.005) - assert_close("dv", ref_dv, tri_dv, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) diff --git a/tests/ops/test_rwkv6.py b/tests/ops/test_rwkv6.py index f379cb8a74..2708f3d90a 100644 --- a/tests/ops/test_rwkv6.py +++ b/tests/ops/test_rwkv6.py @@ -26,13 +26,12 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("head_first", [True, False]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('dtype', [torch.bfloat16]) @pytest.mark.skipif( os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", reason="Skipping test because TEST_CHUNK_VARLEN is enabled" @@ -48,43 +47,38 @@ def test_chunk( D: int, dtype: torch.dtype, gate_logit_normalizer: float, - head_first: bool ): torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' - if head_first: - q = torch.randn((B, H, T, D), dtype=dtype, device=device).requires_grad_() - k = torch.randn((B, H, T, D), dtype=dtype, device=device).requires_grad_() - v = torch.randn((B, H, T, D), dtype=dtype, device=device).requires_grad_() - w = F.logsigmoid(torch.randn((B, H, T, D), dtype=dtype, device=device)) / gate_logit_normalizer - else: - q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() - k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() - v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() - w = F.logsigmoid(torch.randn((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + w = F.logsigmoid(torch.randn((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer u = torch.randn(H, D, dtype=dtype, device=device).requires_grad_(True) h0 = torch.randn(B, H, D, D, dtype=dtype, device=device).requires_grad_() w = w.requires_grad_() do = torch.randn_like(v) - ref, ref_ht = fused_recurrent_rwkv6(q.clone(), - k.clone(), - v.clone(), - w.clone(), - u.clone(), - initial_state=h0.clone(), - output_final_state=True, - head_first=head_first) - ref, _ = fused_recurrent_rwkv6(q.clone(), - k.clone(), - v.clone(), - w.clone(), - u.clone(), - initial_state=h0.clone(), - output_final_state=False, - head_first=head_first) + ref, ref_ht = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + ) + ref, _ = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=False, + ) ((ref * do).sum()).backward() ref_dq, q.grad = q.grad.clone(), None @@ -95,14 +89,15 @@ def test_chunk( ref_dh0, h0.grad = h0.grad.clone(), None # triton implementation - tri, tri_ht = chunk_rwkv6(q.clone(), - k.clone(), - v.clone(), - w.clone(), - u.clone(), - initial_state=h0.clone(), - output_final_state=True, - head_first=head_first) + tri, tri_ht = chunk_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + ) ((tri * do).sum()).backward() tri_dq, q.grad = q.grad.clone(), None tri_dk, k.grad = k.grad.clone(), None @@ -140,7 +135,7 @@ def test_chunk_varlen( torch.manual_seed(42) os.environ['TRITON_F32_DEFAULT'] = 'ieee' # randomly split the sequence into N segments - offsets = torch.cat([ + cu_seqlens = torch.cat([ torch.tensor([0], dtype=torch.long), torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], torch.tensor([T], dtype=torch.long) @@ -162,8 +157,7 @@ def test_chunk_varlen( u.clone(), initial_state=h0.clone(), output_final_state=True, - cu_seqlens=offsets, - head_first=False + cu_seqlens=cu_seqlens, ) ref, _ = fused_recurrent_rwkv6( q.clone(), @@ -173,8 +167,7 @@ def test_chunk_varlen( u.clone(), initial_state=h0.clone(), output_final_state=False, - cu_seqlens=offsets, - head_first=False + cu_seqlens=cu_seqlens, ) ref.backward(do) ref_dq, q.grad = q.grad.clone(), None @@ -192,8 +185,7 @@ def test_chunk_varlen( u.clone(), initial_state=h0.clone(), output_final_state=True, - cu_seqlens=offsets, - head_first=False + cu_seqlens=cu_seqlens, ) tri.backward(do) tri_dq, q.grad = q.grad.clone(), None diff --git a/tests/ops/test_simple_gla.py b/tests/ops/test_simple_gla.py index f2e0bed83e..70794108e2 100644 --- a/tests/ops/test_simple_gla.py +++ b/tests/ops/test_simple_gla.py @@ -108,16 +108,16 @@ def parallel_simple_gla_ref( return o.to(original_dtype), A -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("dtype", [torch.float]) -@pytest.mark.parametrize("scale", [1, 0.1]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('dtype', [torch.float]) +@pytest.mark.parametrize('scale', [1, 0.1]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_chunk( B: int, @@ -165,23 +165,23 @@ def test_chunk( tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_chunk_varlen( N: int, @@ -239,23 +239,23 @@ def test_chunk_varlen( tri_dg, g.grad = g.grad.clone(), None tri_dh0, h0.grad = h0.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" ht", ref_ht, tri_ht, 0.005) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) - assert_close("dh0", ref_dh0, tri_dh0, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' ht', ref_ht, tri_ht, 0.005) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) -@pytest.mark.parametrize("N", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('N', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_parallel_varlen( N: int, @@ -306,23 +306,23 @@ def test_parallel_varlen( tri_dv, v.grad = v.grad.clone(), None tri_dg, g.grad = g.grad.clone(), None - assert_close(" o", ref, tri, 0.004) - assert_close(" dq", ref_dq, tri_dq, 0.005) - assert_close(" dk", ref_dk, tri_dk, 0.005) - assert_close(" dv", ref_dv, tri_dv, 0.005) - assert_close(" dg", ref_dg, tri_dg, 0.005) + assert_close(' o', ref, tri, 0.004) + assert_close(' dq', ref_dq, tri_dq, 0.005) + assert_close(' dk', ref_dk, tri_dk, 0.005) + assert_close(' dv', ref_dv, tri_dv, 0.005) + assert_close(' dg', ref_dg, tri_dg, 0.005) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("gate_logit_normalizer", test_gate_list) -@pytest.mark.parametrize("scale", [0.1]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('gate_logit_normalizer', test_gate_list) +@pytest.mark.parametrize('scale', [0.1]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_parallel( B: int, @@ -358,36 +358,27 @@ def test_parallel( tri_dv, v.grad = v.grad.clone(), None if USE_G: tri_dg, g.grad = g.grad.clone(), None - assert_close(" o", ref, tri, 0.005) - assert_close(" A", ref_A, tri_A, 0.005) - assert_close("dq", ref_dq, tri_dq, 0.005) - assert_close("dk", ref_dk, tri_dk, 0.005) - assert_close("dv", ref_dv, tri_dv, 0.005) + assert_close(' o', ref, tri, 0.005) + assert_close(' A', ref_A, tri_A, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) if USE_G: - assert_close("dg", ref_dg, tri_dg, 0.015) + assert_close('dg', ref_dg, tri_dg, 0.015) -@pytest.mark.parametrize("vary_A", [True, False]) -@pytest.mark.parametrize("dtype", [torch.float, torch.float16]) +@pytest.mark.parametrize('vary_A', [True, False]) +@pytest.mark.parametrize('dtype', [torch.float, torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_simple_gla_to_mamba2(vary_A, dtype): - r""" - Map Mamba-2's `mamba_chunk_scan_combined` kernel to FLA's `simple_gla` kernel - - Dependencies: - $ pip install mamba-ssm==2.2.2 triton==2.3.1 - - Reference: `ssd_minimal_discrete` and `test_correctness` in mamba repository: - https://github.com/state-spaces/mamba/blob/v2.2.2/mamba_ssm/modules/ssd_minimal.py#L82 - """ try: from mamba_ssm.modules.ssd_minimal import ssd_minimal_discrete from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined except ImportError: - pytest.skip("mamba_ssm is not installed.") + pytest.skip('mamba_ssm is not installed.') torch.manual_seed(42) # Dimensions, Denoted (B, T, Q, D, P) in Mamba2 paper @@ -413,11 +404,11 @@ def test_simple_gla_to_mamba2(vary_A, dtype): if not vary_A: # NOTE: fused kernel does not support varying A with time y_fuse, final_fuse = mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=None, return_final_states=True) - assert y_ssd.allclose(y_fuse, 0, atol), f"y diff: {torch.abs(y_ssd - y_fuse).max()}" + assert y_ssd.allclose(y_fuse, 0, atol), f'y diff: {torch.abs(y_ssd - y_fuse).max()}' # fused kernel upcasts state to float32 # https://github.com/state-spaces/mamba/blob/v2.2.2/mamba_ssm/ops/triton/ssd_combined.py#L650 final_fuse = final_fuse.to(dtype) - assert final_ssd.allclose(final_fuse, 0, atol), f"final diff: {torch.abs(final_ssd - final_fuse).max()}" + assert final_ssd.allclose(final_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_fuse).max()}' # mapping inputs Mamba2 -> FLA # C, B, X: [batch, seq, head, hidden] -> [batch, head, seq, hidden] @@ -433,6 +424,6 @@ def test_simple_gla_to_mamba2(vary_A, dtype): # comparing output results between FLA kernel and Mamba2 kernel outputs_gla_fuse, final_gla_fuse = chunk_simple_gla(q, k, v, g, scale=1.0, output_final_state=True) - assert y_rearrange.allclose(outputs_gla_fuse, 0, atol), f"y diff: {torch.abs(y_rearrange - outputs_gla_fuse).max()}" + assert y_rearrange.allclose(outputs_gla_fuse, 0, atol), f'y diff: {torch.abs(y_rearrange - outputs_gla_fuse).max()}' final_gla_fuse = final_gla_fuse.to(dtype) # states hard-coded to float32 in FLA kernel - assert final_rearrange.allclose(final_gla_fuse, 0, atol), f"final diff: {torch.abs(final_ssd - final_gla_fuse).max()}" + assert final_rearrange.allclose(final_gla_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_gla_fuse).max()}' diff --git a/tests/ops/test_solve_tril.py b/tests/ops/test_solve_tril.py index 87c076de83..a15196c4cd 100644 --- a/tests/ops/test_solve_tril.py +++ b/tests/ops/test_solve_tril.py @@ -22,17 +22,17 @@ test_h_list = [2] -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("chunk_size", [16, 32, 64]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('chunk_size', [16, 32, 64]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Pytorch Failure" + reason='Intel Pytorch Failure' ) def test_solve_tril(B, T, H, chunk_size): # do not randomly intiialize A otherwise the inverse is not stable @@ -46,19 +46,19 @@ def test_solve_tril(B, T, H, chunk_size): Ai_ref = torch.inverse(A + torch.eye(A.shape[-1], device=A.device)[None, None, None, ...]) Ai_ref = Ai_ref.reshape(B, H, -1, chunk_size)[:, :, :T, :] - assert_close("solve_tril", Ai, Ai_ref, 0.0001) + assert_close('solve_tril', Ai, Ai_ref, 0.0001) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("cu_seqlens", test_t_varlen_list) -@pytest.mark.parametrize("chunk_size", [64, 32, 16]) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('cu_seqlens', test_t_varlen_list) +@pytest.mark.parametrize('chunk_size', [64, 32, 16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) @pytest.mark.skipif( device_platform == 'intel', - reason="Intel Pytorch Failure" + reason='Intel Pytorch Failure' ) def test_solve_tril_varlen(H, cu_seqlens, chunk_size): T = cu_seqlens[-1] @@ -77,4 +77,4 @@ def test_solve_tril_varlen(H, cu_seqlens, chunk_size): A[:, j:j+actual_size, :, :actual_size].transpose(1, 2) + torch.eye(actual_size, device=A.device, dtype=A.dtype)[None, None, ...] ).transpose(1, 2) - assert_close("solve_tril_varlen", Ai, Ai_ref, 0.0001) + assert_close('solve_tril_varlen', Ai, Ai_ref, 0.0001) diff --git a/tests/ops/test_utils.py b/tests/ops/test_utils.py index 9f2bbfeb6c..1d992806e9 100644 --- a/tests/ops/test_utils.py +++ b/tests/ops/test_utils.py @@ -30,14 +30,14 @@ def reversed_cumsum(x, dim=-1): return y.to(dtype) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_global_cumsum( B: int, @@ -58,14 +58,14 @@ def test_global_cumsum( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_global_cumsum_varlen( B: int, @@ -91,14 +91,14 @@ def test_global_cumsum_varlen( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_global_reversed_cumsum( B: int, @@ -119,14 +119,14 @@ def test_global_reversed_cumsum( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_global_reversed_cumsum_varlen( B: int, @@ -152,12 +152,12 @@ def test_global_reversed_cumsum_varlen( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("C", [32, 64]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('C', [32, 64]) +@pytest.mark.parametrize('dtype', [torch.float16]) def test_local_cumsum( B: int, T: int, @@ -178,15 +178,15 @@ def test_local_cumsum( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_varlen_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("C", [32, 64]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_varlen_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('C', [32, 64]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_local_cumsum_varlen( B: int, @@ -219,15 +219,15 @@ def test_local_cumsum_varlen( torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("C", [32, 64]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('C', [32, 64]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", - reason="Skipping test because TEST_CHUNK_VARLEN is enabled" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled' ) def test_mean_pooling( B: int, @@ -253,15 +253,15 @@ def test_mean_pooling( torch.testing.assert_close(ref_dx, tri_dx.to(ref_dx.dtype), rtol=1.6e-2, atol=3e-5) -@pytest.mark.parametrize("B", test_b_list) -@pytest.mark.parametrize("T", test_t_list) -@pytest.mark.parametrize("H", test_h_list) -@pytest.mark.parametrize("D", test_d_list) -@pytest.mark.parametrize("C", [32, 64]) -@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('D', test_d_list) +@pytest.mark.parametrize('C', [32, 64]) +@pytest.mark.parametrize('dtype', [torch.float16]) @pytest.mark.skipif( - os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", - reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set" + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' ) def test_mean_pooling_varlen( B: int,