diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py index 8edcb8fc88e..b33675519ef 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py @@ -18,15 +18,33 @@ """ from .decode.fused_recurrent import _fused_recurrent_gated_delta_rule_fwd_kernel -from .prefill.chunk import chunk_gated_delta_rule_fwd -from .prefill.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .prefill.chunk_o import chunk_fwd_o +from .prefill.chunk import ( + chunk_gated_delta_rule_fwd, + chunk_gated_delta_rule_fwd_opt, + chunk_gated_delta_rule_fwd_opt_vk, +) +from .prefill.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, + chunk_gated_delta_rule_fwd_h_opt, + chunk_gated_delta_rule_fwd_h_opt_vk, +) +from .prefill.chunk_o import chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk +from .prefill.fused_cumsum_kkt import fused_chunk_local_cumsum_scaled_dot_kkt_fwd +from .prefill.fused_solve_tril_recompute import fused_solve_tril_recompute_w_u from . import gated_delta_rule_utils __all__ = [ "_fused_recurrent_gated_delta_rule_fwd_kernel", "chunk_gated_delta_rule_fwd", + "chunk_gated_delta_rule_fwd_opt", + "chunk_gated_delta_rule_fwd_opt_vk", "chunk_gated_delta_rule_fwd_h", + "chunk_gated_delta_rule_fwd_h_opt", + "chunk_gated_delta_rule_fwd_h_opt_vk", "chunk_fwd_o", + "chunk_fwd_o_opt", + "chunk_fwd_o_opt_vk", + "fused_chunk_local_cumsum_scaled_dot_kkt_fwd", + "fused_solve_tril_recompute_w_u", "gated_delta_rule_utils", ] diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py index ce99ea75e7f..668f971676d 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py @@ -42,6 +42,33 @@ {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} ) +FLA_USE_AUTOTUNE = False + + +def maybe_autotune(configs, default_config=None, **kwargs): + """ + Conditional autotune decorator. + + When FLA_USE_AUTOTUNE is True, behaves identically to @triton.autotune. + When FLA_USE_AUTOTUNE is False (default), uses only the single default_config + (first config in the list if not specified), skipping all benchmark overhead. + + Usage:: + + @maybe_autotune( + configs=[triton.Config(...), triton.Config(...), ...], + default_config=triton.Config({"BV": 64}, num_warps=4, num_stages=2), + key=[...], + ) + @triton.jit + def my_kernel(...): + ... + """ + if FLA_USE_AUTOTUNE: + return triton.autotune(configs=configs, **kwargs) + cfg = default_config if default_config is not None else configs[0] + return triton.autotune(configs=[cfg], **kwargs) + @lru_cache(maxsize=1) def check_environments(): diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py index 31021b1f850..5347263c9f3 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py @@ -8,16 +8,36 @@ This module provides optimized Triton kernels for prefill/training operations. """ -from .chunk import chunk_gated_delta_rule_fwd -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .chunk_o import chunk_fwd_o -from .fused_cumsum_kkt import fused_cumsum_kkt +from .chunk import ( + chunk_gated_delta_rule_fwd, + chunk_gated_delta_rule_fwd_opt, + chunk_gated_delta_rule_fwd_opt_vk, +) +from .chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, + chunk_gated_delta_rule_fwd_h_opt, + chunk_gated_delta_rule_fwd_h_opt_vk, +) +from .chunk_o import chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk +from .fused_cumsum_kkt import ( + fused_cumsum_kkt, + fused_chunk_local_cumsum_scaled_dot_kkt_fwd, +) +from .fused_solve_tril_recompute import fused_solve_tril_recompute_w_u from .fused_gdn_gating_prefill import fused_gdn_gating_and_sigmoid __all__ = [ "chunk_gated_delta_rule_fwd", + "chunk_gated_delta_rule_fwd_opt", + "chunk_gated_delta_rule_fwd_opt_vk", "chunk_gated_delta_rule_fwd_h", + "chunk_gated_delta_rule_fwd_h_opt", + "chunk_gated_delta_rule_fwd_h_opt_vk", "chunk_fwd_o", + "chunk_fwd_o_opt", + "chunk_fwd_o_opt_vk", "fused_cumsum_kkt", + "fused_chunk_local_cumsum_scaled_dot_kkt_fwd", + "fused_solve_tril_recompute_w_u", "fused_gdn_gating_and_sigmoid", ] diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py index 9713b904956..974b207b6fd 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py @@ -11,8 +11,14 @@ import torch -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .chunk_o import chunk_fwd_o +from .chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, + chunk_gated_delta_rule_fwd_h_opt, + chunk_gated_delta_rule_fwd_h_opt_vk, +) +from .chunk_o import chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk +from .fused_cumsum_kkt import fused_chunk_local_cumsum_scaled_dot_kkt_fwd +from .fused_solve_tril_recompute import fused_solve_tril_recompute_w_u from ..utils import ( chunk_local_cumsum, chunk_scaled_dot_kkt_fwd, @@ -106,3 +112,157 @@ def chunk_gated_delta_rule_fwd( ) return g, o, A, final_state + + +def chunk_gated_delta_rule_fwd_opt( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, +): + """ + Optimized chunk gated delta rule forward computation (Forward only). + + This function implements an optimized chunk-based parallel computation for + the gated delta rule, using fused kernels and transposed intermediate layouts + to reduce global memory round-trips. + + Note: This implementation only supports forward pass. Backward pass is not available. + + Args: + q: Query tensor of shape [B, T, Hg, K] + k: Key tensor of shape [B, T, Hg, K] + v: Value tensor of shape [B, T, H, V] + g: Gate tensor (in log space, pre-cumsum) of shape [B, T, H] + beta: Beta parameter tensor of shape [B, T, H] + scale: Scaling factor for queries + initial_state: Initial hidden state of shape [N, H, K, V] + output_final_state: Whether to output the final state + cu_seqlens: Cumulative sequence lengths for variable-length inputs (optional) [N+1] + + Returns: + tuple: (g_cumsum, o, final_state) where: + - g_cumsum: Cumulative gate values [B, T, H] + - o: Output tensor [B, T, H, V] + - final_state: Final hidden state [N, H, K, V] if output_final_state=True, else None + """ + # Step 1: Compute fused local cumulative sum of gates and KKT + g_cumsum, A_raw = fused_chunk_local_cumsum_scaled_dot_kkt_fwd( + k=k, + beta=beta, + g=g, + cu_seqlens=cu_seqlens, + ) + + # Step 2: Compute fused triangular solve and recompute w, u + # w, u are already in [B, H, T, K/V] head-major contiguous layout + w, u = fused_solve_tril_recompute_w_u( + A_raw=A_raw, + k=k, + v=v, + beta=beta, + g_cumsum=g_cumsum, + cu_seqlens=cu_seqlens, + ) + + # Step 3: Compute hidden states + h, v_new, final_state = chunk_gated_delta_rule_fwd_h_opt( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + + # Step 4: Compute output (directly in [B, T, H, V] layout) + o = chunk_fwd_o_opt( + q=q, + k=k, + v=v_new, + h=h, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + ) + + return g_cumsum, o, final_state + + +def chunk_gated_delta_rule_fwd_opt_vk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, +): + """ + Optimized chunk gated delta rule forward with h layout [V, K]. + + Uses the same fused K12/K34 kernels as opt, but K5/K6 use transposed + h layout [V, K] instead of [K, V]. + + Args: + q: [B, T, Hg, K] + k: [B, T, Hg, K] + v: [B, T, H, V] + g: [B, T, H] — raw gate (pre-cumsum) + beta: [B, T, H] + scale: float + initial_state: [N, H, V, K] — note transposed h layout + output_final_state: bool + cu_seqlens: [N+1] optional + + Returns: + tuple: (g_cumsum, o, final_state) where: + - g_cumsum: [B, T, H] + - o: [B, T, H, V] + - final_state: [N, H, V, K] if output_final_state=True, else None + """ + g_cumsum, A_raw = fused_chunk_local_cumsum_scaled_dot_kkt_fwd( + k=k, + beta=beta, + g=g, + cu_seqlens=cu_seqlens, + ) + + w, u = fused_solve_tril_recompute_w_u( + A_raw=A_raw, + k=k, + v=v, + beta=beta, + g_cumsum=g_cumsum, + cu_seqlens=cu_seqlens, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h_opt_vk( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + + o = chunk_fwd_o_opt_vk( + q=q, + k=k, + v=v_new, + h=h, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + ) + + return g_cumsum, o, final_state diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py index 8f85880efa2..0c18f11ffe2 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py @@ -22,6 +22,7 @@ USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem, + maybe_autotune, ) NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] @@ -39,7 +40,7 @@ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] @@ -295,7 +296,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] @@ -643,6 +644,680 @@ def grid(meta): return h, v_new, final_state +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] 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, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@maybe_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in NUM_STAGES_FWD + for BV in [16, 32, 64] + ], + key=["H", "K", "V", "BT", "IS_VARLEN"], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T", "T_flat"]) +def chunk_gated_delta_rule_fwd_kernel_h_opt( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + T_flat, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: 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(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) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + h += ((boh * H + i_h) * K * V).to(tl.int64) + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) + if IS_VARLEN: + v += ((i_h * T_flat + bos) * V).to(tl.int64) + w += ((i_h * T_flat + bos) * K).to(tl.int64) + else: + v += (((i_n * H + i_h) * T_flat) * V).to(tl.int64) + w += (((i_n * H + i_h) * T_flat) * K).to(tl.int64) + stride_v = V + stride_w = K + if SAVE_NEW_VALUE: + if IS_VARLEN: + v_new += ((i_h * T_flat + bos) * V).to(tl.int64) + else: + v_new += (((i_n * H + i_h) * T_flat) * V).to(tl.int64) + stride_h = H * K * V + stride_k = Hg * K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr( + h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr( + h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr( + h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + p_h1 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_v = tl.make_block_ptr( + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + if SAVE_NEW_VALUE: + p_vn = tl.make_block_ptr( + v_new, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_vn, b_v.to(p_vn.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k1, + mask=(o_k1 < K), + other=0.0, + ) + b_h1 *= exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k2, + mask=(o_k2 < K), + other=0.0, + ) + b_h2 *= exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k3, + mask=(o_k3 < K), + other=0.0, + ) + b_h3 *= exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k4, + mask=(o_k4 < K), + other=0.0, + ) + b_h4 *= exp(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_h_opt( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """ + Optimized hidden state forward with Hg-aware k strides. + + w and u are expected in head-major contiguous layout [B, H, T, K] / [B, H, T, V]. + v_new output is [B, H, T_flat, V]. + """ + B, T, Hg, K = k.shape + BT = chunk_size + + H = w.shape[1] + V = u.shape[-1] + T_flat = w.shape[2] + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + N = len(cu_seqlens) - 1 + NT = len(chunk_indices) + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + else: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + + assert K <= 256, "current kernel does not support head dimension larger than 256." + + h = k.new_empty(B, NT, H, K, V) + final_state = ( + k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + ) + v_new = k.new_empty(B, H, T_flat, V, dtype=u.dtype) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + + chunk_gated_delta_rule_fwd_kernel_h_opt[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + T_flat=T_flat, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return h, v_new, final_state + + +# ===================================================================== +# opt_vk variant: h layout [V, K] (transposed from opt's [K, V]) +# All other layouts (k, w, u, v_new) are identical to opt. +# ===================================================================== + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] 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, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@maybe_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in NUM_STAGES_FWD + for BV in [16, 32, 64] + ], + key=["H", "K", "V", "BT"], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T", "T_flat"]) +def chunk_gated_delta_rule_fwd_kernel_h_opt_vk( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + T_flat, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: 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(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) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BV, 64] — h in [V, K] layout (transposed from opt's [64, BV]) + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + + h += ((boh * H + i_h) * V * K).to(tl.int64) + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) + if IS_VARLEN: + v += ((i_h * T_flat + bos) * V).to(tl.int64) + w += ((i_h * T_flat + bos) * K).to(tl.int64) + else: + v += (((i_n * H + i_h) * T_flat) * V).to(tl.int64) + w += (((i_n * H + i_h) * T_flat) * K).to(tl.int64) + stride_v = V + stride_w = K + if SAVE_NEW_VALUE: + if IS_VARLEN: + v_new += ((i_h * T_flat + bos) * V).to(tl.int64) + else: + v_new += (((i_n * H + i_h) * T_flat) * V).to(tl.int64) + stride_h = H * V * K + stride_k = Hg * K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * V * K + if STORE_FINAL_STATE: + ht = ht + i_nh * V * K + + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + # Store h snapshot [V, K] + p_h1 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + # b_v = u - w @ h^T (h is [BV,64], need [64,BV] for dot with w[BT,64]) + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + p_v = tl.make_block_ptr( + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + if SAVE_NEW_VALUE: + p_vn = tl.make_block_ptr( + v_new, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_vn, b_v.to(p_vn.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k1, + mask=(o_k1 < K), + other=0.0, + ) + b_h1 *= exp(b_gk_last1)[None, :] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k2, + mask=(o_k2 < K), + other=0.0, + ) + b_h2 *= exp(b_gk_last2)[None, :] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k3, + mask=(o_k3 < K), + other=0.0, + ) + b_h3 *= exp(b_gk_last3)[None, :] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k4, + mask=(o_k4 < K), + other=0.0, + ) + b_h4 *= exp(b_gk_last4)[None, :] + b_v = b_v.to(k.dtype.element_ty) + + # h[V,K] += v_new^T @ k → [BV,64] += trans(dot(k[64,BT], v[BT,BV])) + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.trans(tl.dot(b_k, b_v)) + if K > 64: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.trans(tl.dot(b_k, b_v)) + if K > 128: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.trans(tl.dot(b_k, b_v)) + if K > 192: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.trans(tl.dot(b_k, b_v)) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_h_opt_vk( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """ + Optimized hidden state forward with h layout [V, K]. + + w and u are expected in head-major contiguous layout [B, H, T, K] / [B, H, T, V]. + initial_state/final_state: [N, H, V, K]. + h snapshots: [B, NT, H, V, K]. + v_new output is [B, H, T_flat, V]. + """ + B, T, Hg, K = k.shape + BT = chunk_size + + H = w.shape[1] + V = u.shape[-1] + T_flat = w.shape[2] + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + N = len(cu_seqlens) - 1 + NT = len(chunk_indices) + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + else: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + + assert K <= 256, "current kernel does not support head dimension larger than 256." + + h = k.new_empty(B, NT, H, V, K) + final_state = ( + k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None + ) + v_new = k.new_empty(B, H, T_flat, V, dtype=u.dtype) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + + chunk_gated_delta_rule_fwd_kernel_h_opt_vk[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + T_flat=T_flat, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return h, v_new, final_state + + def chunk_gated_delta_rule_bwd_dhu( q: torch.Tensor, k: torch.Tensor, diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py index ad81a59d539..93b47266fe8 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py @@ -16,6 +16,7 @@ IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem, + maybe_autotune, ) from ..utils import prepare_chunk_indices from ..utils.op import exp @@ -31,7 +32,7 @@ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({"BK": 128, "BV": 128}, num_warps=8, num_stages=3), triton.Config({"BK": 64, "BV": 64}, num_warps=4, num_stages=3), @@ -153,7 +154,7 @@ def chunk_fwd_kernel_o( "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS @@ -360,7 +361,7 @@ def chunk_bwd_kernel_dqkwg( "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS @@ -474,7 +475,7 @@ def chunk_bwd_kernel_dv( "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } ) -@triton.autotune( +@maybe_autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS @@ -622,6 +623,364 @@ def grid(meta): return o +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@maybe_autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "IS_VARLEN"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T", "T_flat"]) +def chunk_fwd_kernel_o_opt( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + T_flat, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, 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_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos = i_b * T + + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + if IS_VARLEN: + v += ((i_h * T_flat + bos) * V).to(tl.int64) + o += ((bos * H + i_h) * V).to(tl.int64) + else: + v += (((i_b * H + i_h) * T_flat) * V).to(tl.int64) + o += ((i_b * T * H + i_h) * V).to(tl.int64) + h += (i_tg * H + i_h).to(tl.int64) * K * V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + + b_o += tl.dot(b_q, b_h) + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_o_opt( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +) -> torch.Tensor: + """ + Optimized output forward with transposed v layout and Hg-aware q/k strides. + + Args: + q: [B, T, Hg, K] + k: [B, T, Hg, K] + v: [B, H, T, V] + h: [B, NT, H, K, V] + g: [B*T, H] FP32 + scale: float + cu_seqlens: [N+1] + chunk_size: int + + Returns: + o: [B, T, H, V] + """ + B, T, Hg, K = q.shape + H = v.shape[1] + T_flat = v.shape[2] + V = v.shape[-1] + BT = chunk_size + 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 + + o = v.new_empty(B, T, H, V) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), NT, B * H) + + chunk_fwd_kernel_o_opt[grid]( + q=q, + k=k, + v=v, + h=h, + g=g, + o=o, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + T_flat=T_flat, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return o + + +# ===================================================================== +# opt_vk variant: h layout [V, K] (transposed from opt's [K, V]) +# All other layouts identical to opt. +# ===================================================================== + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@maybe_autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T", "T_flat"]) +def chunk_fwd_kernel_o_opt_vk( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + T_flat, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, 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_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos = i_b * T + + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + if IS_VARLEN: + v += ((i_h * T_flat + bos) * V).to(tl.int64) + o += ((bos * H + i_h) * V).to(tl.int64) + else: + v += (((i_b * H + i_h) * T_flat) * V).to(tl.int64) + o += ((i_b * T * H + i_h) * V).to(tl.int64) + h += (i_tg * H + i_h).to(tl.int64) * V * K + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + + b_o += tl.dot(b_q, tl.trans(b_h)) + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_o_opt_vk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +) -> torch.Tensor: + """ + Optimized output forward with h layout [V, K]. + + Args: + q: [B, T, Hg, K] + k: [B, T, Hg, K] + v: [B, H, T, V] (token-major from K5 opt_vk) + h: [B, NT, H, V, K] (h layout [V, K]) + g: [B*T, H] FP32 + scale: float + cu_seqlens: [N+1] + chunk_size: int + + Returns: + o: [B, T, H, V] + """ + B, T, Hg, K = q.shape + H = v.shape[1] + T_flat = v.shape[2] + V = v.shape[-1] + BT = chunk_size + 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 + + o = v.new_empty(B, T, H, V) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), NT, B * H) + + chunk_fwd_kernel_o_opt_vk[grid]( + q=q, + k=k, + v=v, + h=h, + g=g, + o=o, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + T_flat=T_flat, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return o + + def chunk_bwd_dv( q: torch.Tensor, k: torch.Tensor, diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py index 276deda595b..c4217f74a31 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py @@ -2,7 +2,9 @@ import triton import triton.language as tl +from ..gated_delta_rule_utils import autotune_cache_kwargs, IS_AMD, maybe_autotune from ..utils import prepare_chunk_indices +from ..utils.op import exp @triton.jit @@ -133,3 +135,148 @@ def fused_cumsum_kkt( num_stages=3, ) return g_cumsum, A + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@maybe_autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=nw, num_stages=ns) + for BK in [32, 64] + for nw in [2, 4] + for ns in ([2, 3] if IS_AMD else [2, 3, 4]) + ], + key=["H", "K", "BT", "IS_VARLEN"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def fused_chunk_local_cumsum_scaled_dot_kkt_fwd_kernel( + g, + k, + beta, + g_cumsum_out, + A_out, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + 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(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 = i_b * T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g_cumsum = tl.cumsum(b_g, axis=0) + + p_go = tl.make_block_ptr( + g_cumsum_out + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + tl.store(p_go, b_g_cumsum.to(p_go.dtype.element_ty), boundary_check=(0,)) + + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_beta[:, None] + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k)) + + b_g_diff = b_g_cumsum[:, None] - b_g_cumsum[None, :] + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A * exp(b_g_diff), 0.0) + + p_A = tl.make_block_ptr( + A_out + (bos * H + i_h) * BT, + (T, BT), + (BT * H, 1), + (i_t * BT, 0), + (BT, BT), + (1, 0), + ) + tl.store(p_A, b_A.to(A_out.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_chunk_local_cumsum_scaled_dot_kkt_fwd( + k: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + g_output_dtype: torch.dtype = torch.float32, + A_output_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused cumsum + scaled dot KKT (optimized, with autotuning). + + Args: + k: [B, T, Hg, K] + beta: [B, T, H] + g: [B, T, H], raw forget gate increments + cu_seqlens: [N+1] + chunk_size: int (must be 64) + g_output_dtype: dtype for g_cumsum (default fp32) + A_output_dtype: dtype for A_raw (default fp32) + + Returns: + g_cumsum: [B, T, H] + A_raw: [B, T, H, 64] + """ + B, T, Hg, K = k.shape + H = beta.shape[-1] + BT = chunk_size + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) + else: + chunk_indices = None + NT = triton.cdiv(T, BT) + + g_cumsum_out = torch.empty(B, T, H, device=g.device, dtype=g_output_dtype) + A_out = torch.empty(B, T, H, BT, device=k.device, dtype=A_output_dtype) + + fused_chunk_local_cumsum_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( + g, + k, + beta, + g_cumsum_out, + A_out, + cu_seqlens, + chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + BT=BT, + ) + return g_cumsum_out, A_out diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py new file mode 100644 index 00000000000..74aa2656a6c --- /dev/null +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# Adapted from flash-linear-attention: Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +Fused triangular solve + recompute w, u in a single kernel. + +Eliminates the intermediate Ai tensor (64x64 per chunk x head) global +memory round-trip by keeping the inverse blocks in registers. +""" + +import torch +import triton +import triton.language as tl + +from ..gated_delta_rule_utils import autotune_cache_kwargs, IS_AMD, maybe_autotune +from ..utils import prepare_chunk_indices +from ..utils.op import exp +from ..utils.solve_tril import FLA_TRIL_PRECISION + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@maybe_autotune( + configs=[ + triton.Config({}, num_warps=nw, num_stages=ns) + for nw in [2, 4, 8] + for ns in ([2, 3] if IS_AMD else [2, 3, 4]) + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def fused_solve_tril_recompute_w_u_kernel( + A_raw, + k, + v, + beta, + g, + w, + u, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + T_flat = T + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + # ================================================================ + # Phase 1: compute (I + A)^{-1} in registers (triangular solve) + # ================================================================ + o_i = tl.arange(0, 16) + m_lo = o_i[:, None] > o_i[None, :] + m_id = o_i[:, None] == o_i[None, :] + A_base = A_raw + (bos * H + i_h) * BT + + p11 = tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p22 = tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + p33 = tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + ) + p44 = tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + ) + b11 = -tl.where(m_lo, tl.load(p11, boundary_check=(0, 1)).to(tl.float32), 0) + b22 = -tl.where(m_lo, tl.load(p22, boundary_check=(0, 1)).to(tl.float32), 0) + b33 = -tl.where(m_lo, tl.load(p33, boundary_check=(0, 1)).to(tl.float32), 0) + b44 = -tl.where(m_lo, tl.load(p44, boundary_check=(0, 1)).to(tl.float32), 0) + + for i in range(2, min(16, T - i_t * BT)): + r = -tl.load(A_base + (i_t * BT + i) * H * BT + o_i) + r = r + tl.sum(r[:, None] * b11, 0) + b11 = tl.where((o_i == i)[:, None], r, b11) + for i in range(18, min(32, T - i_t * BT)): + r = -tl.load(A_base + (i_t * BT + i) * H * BT + o_i + 16) + r = r + tl.sum(r[:, None] * b22, 0) + b22 = tl.where((o_i == i - 16)[:, None], r, b22) + for i in range(34, min(48, T - i_t * BT)): + r = -tl.load(A_base + (i_t * BT + i) * H * BT + o_i + 32) + r = r + tl.sum(r[:, None] * b33, 0) + b33 = tl.where((o_i == i - 32)[:, None], r, b33) + for i in range(50, min(64, T - i_t * BT)): + r = -tl.load(A_base + (i_t * BT + i) * H * BT + o_i + 48) + r = r + tl.sum(r[:, None] * b44, 0) + b44 = tl.where((o_i == i - 48)[:, None], r, b44) + b11 += m_id + b22 += m_id + b33 += m_id + b44 += m_id + + rA21 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + rA31 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + rA32 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + rA41 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + rA42 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + rA43 = tl.load( + tl.make_block_ptr( + A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) + ), + boundary_check=(0, 1), + ).to(tl.float32) + + b21 = -tl.dot( + tl.dot(b22, rA21, input_precision=DOT_PRECISION), + b11, + input_precision=DOT_PRECISION, + ) + b32 = -tl.dot( + tl.dot(b33, rA32, input_precision=DOT_PRECISION), + b22, + input_precision=DOT_PRECISION, + ) + b43 = -tl.dot( + tl.dot(b44, rA43, input_precision=DOT_PRECISION), + b33, + input_precision=DOT_PRECISION, + ) + b31 = -tl.dot( + b33, + tl.dot(rA31, b11, input_precision=DOT_PRECISION) + + tl.dot(rA32, b21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b42 = -tl.dot( + b44, + tl.dot(rA42, b22, input_precision=DOT_PRECISION) + + tl.dot(rA43, b32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b41 = -tl.dot( + b44, + tl.dot(rA41, b11, input_precision=DOT_PRECISION) + + tl.dot(rA42, b21, input_precision=DOT_PRECISION) + + tl.dot(rA43, b31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + h11 = b11.to(tl.bfloat16) + h22 = b22.to(tl.bfloat16) + h33 = b33.to(tl.bfloat16) + h44 = b44.to(tl.bfloat16) + h21 = b21.to(tl.bfloat16) + h31 = b31.to(tl.bfloat16) + h32 = b32.to(tl.bfloat16) + h41 = b41.to(tl.bfloat16) + h42 = b42.to(tl.bfloat16) + h43 = b43.to(tl.bfloat16) + + # ================================================================ + # Phase 2: u = Ai @ (v * beta), w = Ai @ (k * beta * exp(g)) + # ================================================================ + beta_base = beta + bos * H + i_h + g_base = g + bos * H + i_h + + p_b0 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT,), (16,), (0,)) + p_b1 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 16,), (16,), (0,)) + p_b2 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 32,), (16,), (0,)) + p_b3 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 48,), (16,), (0,)) + bb0 = tl.load(p_b0, boundary_check=(0,)) + bb1 = tl.load(p_b1, boundary_check=(0,)) + bb2 = tl.load(p_b2, boundary_check=(0,)) + bb3 = tl.load(p_b3, boundary_check=(0,)) + + p_g0 = tl.make_block_ptr(g_base, (T,), (H,), (i_t * BT,), (16,), (0,)) + p_g1 = tl.make_block_ptr(g_base, (T,), (H,), (i_t * BT + 16,), (16,), (0,)) + p_g2 = tl.make_block_ptr(g_base, (T,), (H,), (i_t * BT + 32,), (16,), (0,)) + p_g3 = tl.make_block_ptr(g_base, (T,), (H,), (i_t * BT + 48,), (16,), (0,)) + eg0 = exp(tl.load(p_g0, boundary_check=(0,))) + eg1 = exp(tl.load(p_g1, boundary_check=(0,))) + eg2 = exp(tl.load(p_g2, boundary_check=(0,))) + eg3 = exp(tl.load(p_g3, boundary_check=(0,))) + + v_base = v + (bos * H + i_h) * V + if IS_VARLEN: + u_base = u + (i_h * T_flat + bos) * V + else: + u_base = u + (((i_b * H + i_h) * T_flat) * V) + + for i_v in range(tl.cdiv(V, BV)): + pv0 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (16, BV), (1, 0) + ) + pv1 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_t * BT + 16, i_v * BV), (16, BV), (1, 0) + ) + pv2 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_t * BT + 32, i_v * BV), (16, BV), (1, 0) + ) + pv3 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_t * BT + 48, i_v * BV), (16, BV), (1, 0) + ) + vb0 = (tl.load(pv0, boundary_check=(0, 1)) * bb0[:, None]).to(tl.bfloat16) + vb1 = (tl.load(pv1, boundary_check=(0, 1)) * bb1[:, None]).to(tl.bfloat16) + vb2 = (tl.load(pv2, boundary_check=(0, 1)) * bb2[:, None]).to(tl.bfloat16) + vb3 = (tl.load(pv3, boundary_check=(0, 1)) * bb3[:, None]).to(tl.bfloat16) + + u0 = tl.dot(h11, vb0, allow_tf32=False) + u1 = tl.dot(h21, vb0, allow_tf32=False) + tl.dot(h22, vb1, allow_tf32=False) + u2 = ( + tl.dot(h31, vb0, allow_tf32=False) + + tl.dot(h32, vb1, allow_tf32=False) + + tl.dot(h33, vb2, allow_tf32=False) + ) + u3 = ( + tl.dot(h41, vb0, allow_tf32=False) + + tl.dot(h42, vb1, allow_tf32=False) + + tl.dot(h43, vb2, allow_tf32=False) + + tl.dot(h44, vb3, allow_tf32=False) + ) + + pu0 = tl.make_block_ptr( + u_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (16, BV), (1, 0) + ) + pu1 = tl.make_block_ptr( + u_base, (T, V), (V, 1), (i_t * BT + 16, i_v * BV), (16, BV), (1, 0) + ) + pu2 = tl.make_block_ptr( + u_base, (T, V), (V, 1), (i_t * BT + 32, i_v * BV), (16, BV), (1, 0) + ) + pu3 = tl.make_block_ptr( + u_base, (T, V), (V, 1), (i_t * BT + 48, i_v * BV), (16, BV), (1, 0) + ) + tl.store(pu0, u0.to(pu0.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pu1, u1.to(pu1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pu2, u2.to(pu2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pu3, u3.to(pu3.dtype.element_ty), boundary_check=(0, 1)) + + k_base = k + (bos * Hg + i_h // (H // Hg)) * K + if IS_VARLEN: + w_base = w + (i_h * T_flat + bos) * K + else: + w_base = w + (((i_b * H + i_h) * T_flat) * K) + + for i_k in range(tl.cdiv(K, BK)): + pk0 = tl.make_block_ptr( + k_base, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (16, BK), (1, 0) + ) + pk1 = tl.make_block_ptr( + k_base, (T, K), (Hg * K, 1), (i_t * BT + 16, i_k * BK), (16, BK), (1, 0) + ) + pk2 = tl.make_block_ptr( + k_base, (T, K), (Hg * K, 1), (i_t * BT + 32, i_k * BK), (16, BK), (1, 0) + ) + pk3 = tl.make_block_ptr( + k_base, (T, K), (Hg * K, 1), (i_t * BT + 48, i_k * BK), (16, BK), (1, 0) + ) + kb0 = (tl.load(pk0, boundary_check=(0, 1)) * bb0[:, None] * eg0[:, None]).to( + tl.bfloat16 + ) + kb1 = (tl.load(pk1, boundary_check=(0, 1)) * bb1[:, None] * eg1[:, None]).to( + tl.bfloat16 + ) + kb2 = (tl.load(pk2, boundary_check=(0, 1)) * bb2[:, None] * eg2[:, None]).to( + tl.bfloat16 + ) + kb3 = (tl.load(pk3, boundary_check=(0, 1)) * bb3[:, None] * eg3[:, None]).to( + tl.bfloat16 + ) + + w0 = tl.dot(h11, kb0) + w1 = tl.dot(h21, kb0) + tl.dot(h22, kb1) + w2 = tl.dot(h31, kb0) + tl.dot(h32, kb1) + tl.dot(h33, kb2) + w3 = tl.dot(h41, kb0) + tl.dot(h42, kb1) + tl.dot(h43, kb2) + tl.dot(h44, kb3) + + pw0 = tl.make_block_ptr( + w_base, (T, K), (K, 1), (i_t * BT, i_k * BK), (16, BK), (1, 0) + ) + pw1 = tl.make_block_ptr( + w_base, (T, K), (K, 1), (i_t * BT + 16, i_k * BK), (16, BK), (1, 0) + ) + pw2 = tl.make_block_ptr( + w_base, (T, K), (K, 1), (i_t * BT + 32, i_k * BK), (16, BK), (1, 0) + ) + pw3 = tl.make_block_ptr( + w_base, (T, K), (K, 1), (i_t * BT + 48, i_k * BK), (16, BK), (1, 0) + ) + tl.store(pw0, w0.to(pw0.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pw1, w1.to(pw1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pw2, w2.to(pw2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(pw3, w3.to(pw3.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_solve_tril_recompute_w_u( + A_raw: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g_cumsum: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused triangular solve + recompute w, u in a single kernel. + + Args: + A_raw: [B, T, H, BT=64], strictly lower triangular + k: [B, T, Hg, K] + v: [B, T, H, V] + beta: [B, T, H] + g_cumsum: [B, T, H] FP32, cumulative gate + cu_seqlens: [N+1] + + Returns: + w: [B, H, T, K], head-major contiguous layout + u: [B, H, T, V], head-major contiguous layout + """ + B, T, Hg, K, V = *k.shape, v.shape[-1] + H = v.shape[-2] + BT = A_raw.shape[-1] + BK = 64 + BV = 64 + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) + else: + chunk_indices = None + NT = triton.cdiv(T, BT) + + u_out = v.new_empty(B, H, T, V) + w_out = k.new_empty(B, H, T, K) + + fused_solve_tril_recompute_w_u_kernel[(NT, B * H)]( + A_raw, + k, + v, + beta, + g_cumsum, + w_out, + u_out, + cu_seqlens, + chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION=FLA_TRIL_PRECISION, + ) + return w_out, u_out diff --git a/aiter/ops/triton/gated_delta_net/__init__.py b/aiter/ops/triton/gated_delta_net/__init__.py index bc46e85b4d0..066c0d063eb 100644 --- a/aiter/ops/triton/gated_delta_net/__init__.py +++ b/aiter/ops/triton/gated_delta_net/__init__.py @@ -10,10 +10,14 @@ from .gated_delta_rule import ( chunk_gated_delta_rule, + chunk_gated_delta_rule_opt, + chunk_gated_delta_rule_opt_vk, fused_recurrent_gated_delta_rule, ) __all__ = [ "fused_recurrent_gated_delta_rule", "chunk_gated_delta_rule", + "chunk_gated_delta_rule_opt", + "chunk_gated_delta_rule_opt_vk", ] diff --git a/aiter/ops/triton/gated_delta_net/gated_delta_rule.py b/aiter/ops/triton/gated_delta_net/gated_delta_rule.py index 7a15c38340d..4615067ed4f 100644 --- a/aiter/ops/triton/gated_delta_net/gated_delta_rule.py +++ b/aiter/ops/triton/gated_delta_net/gated_delta_rule.py @@ -22,6 +22,8 @@ from aiter.ops.triton._triton_kernels.gated_delta_rule import ( _fused_recurrent_gated_delta_rule_fwd_kernel, chunk_gated_delta_rule_fwd, + chunk_gated_delta_rule_fwd_opt, + chunk_gated_delta_rule_fwd_opt_vk, ) from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( l2norm_fwd, @@ -333,3 +335,179 @@ def chunk_gated_delta_rule( cu_seqlens=cu_seqlens, ) return o.to(q.dtype), final_state + + +def chunk_gated_delta_rule_opt( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Optimized chunk-based gated delta rule operation using Triton (Forward only). + + This function implements an optimized chunk-based parallel computation for the + gated delta rule, using fused kernels and transposed intermediate layouts to + reduce global memory round-trips. + + Warning: + This function only supports forward pass and does NOT compute gradients. + Do not use this for training. For training, use flash-linear-attention library. + + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + g (decays in log space) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (float, optional): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (torch.Tensor, optional): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (bool): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to use L2 normalization in the kernel. Default: `False`. + cu_seqlens (torch.LongTensor, optional): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. Default: `None`. + + Returns: + tuple[torch.Tensor, torch.Tensor]: + - o (torch.Tensor): Outputs of shape `[B, T, H, V]`. + - final_state (torch.Tensor): Final state of shape `[N, H, K, V]` if + `output_final_state=True` else `None`. + + Raises: + ValueError: If input shapes are invalid when using cu_seqlens. + """ + # Input validation + 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 `cu_seqlens`. " + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + + # Set default values + if scale is None: + scale = k.shape[-1] ** -0.5 + + # Log operation + _LOGGER.info( + f"CHUNK_GATED_DELTA_RULE_OPT: q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"scale={scale}, use_qk_l2norm={use_qk_l2norm_in_kernel}" + ) + + # Apply L2 normalization if requested + if use_qk_l2norm_in_kernel: + _LOGGER.info("Applying L2 normalization to q and k") + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + # Call aiter's optimized chunk forward pass + g_cumsum, o, final_state = chunk_gated_delta_rule_fwd_opt( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o.to(q.dtype), final_state + + +def chunk_gated_delta_rule_opt_vk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Optimized chunk-based gated delta rule with h layout [V, K] (Forward only). + + Same fused K12/K34 kernels as chunk_gated_delta_rule_opt, but K5/K6 + use transposed hidden state layout [V, K] instead of [K, V]. + + Args: + q (torch.Tensor): queries of shape `[B, T, H, K]`. + k (torch.Tensor): keys of shape `[B, T, H, K]`. + v (torch.Tensor): values of shape `[B, T, H, V]`. + g (torch.Tensor): g (decays in log space) of shape `[B, T, H]`. + beta (torch.Tensor): betas of shape `[B, T, H]`. + scale (float, optional): Scale factor. Default: `1 / sqrt(K)`. + initial_state (torch.Tensor, optional): + Initial state of shape `[N, H, V, K]` — note transposed layout. + output_final_state (bool): Whether to output final state `[N, H, V, K]`. + use_qk_l2norm_in_kernel (bool): Whether to use L2 normalization. + cu_seqlens (torch.LongTensor, optional): Cumulative sequence lengths `[N+1]`. + + Returns: + tuple[torch.Tensor, torch.Tensor]: + - o: Outputs of shape `[B, T, H, V]`. + - final_state: `[N, H, V, K]` if `output_final_state=True` else `None`. + """ + 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 `cu_seqlens`." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + + if scale is None: + scale = k.shape[-1] ** -0.5 + + _LOGGER.info( + f"CHUNK_GATED_DELTA_RULE_OPT_VK: q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"scale={scale}, use_qk_l2norm={use_qk_l2norm_in_kernel}" + ) + + if use_qk_l2norm_in_kernel: + _LOGGER.info("Applying L2 normalization to q and k") + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + g_cumsum, o, final_state = chunk_gated_delta_rule_fwd_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o.to(q.dtype), final_state diff --git a/op_tests/triton_tests/test_gated_delta_rule.py b/op_tests/triton_tests/test_gated_delta_rule.py index 5bc43f56af7..c6e59ebadc1 100644 --- a/op_tests/triton_tests/test_gated_delta_rule.py +++ b/op_tests/triton_tests/test_gated_delta_rule.py @@ -10,6 +10,8 @@ from aiter.ops.triton.gated_delta_net import ( fused_recurrent_gated_delta_rule, chunk_gated_delta_rule, + chunk_gated_delta_rule_opt, + chunk_gated_delta_rule_opt_vk, ) from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, @@ -535,3 +537,369 @@ def test_chunk_varlen( # 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) + + +@pytest.mark.parametrize( + ( + "B", + "T", + "H", + "D", + "scale", + "gate_logit_normalizer", + "mask_p", + "use_qk_l2norm_in_kernel", + "dtype", + ), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format( + *test + ), + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16), + # bfloat16 tests + (2, 500, 3, 60, 1, 1, 0, False, torch.bfloat16), + (2, 1000, 3, 64, 0.1, 1, 0, False, torch.bfloat16), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.bfloat16), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.bfloat16), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.bfloat16), + ] + ], +) +def test_chunk_opt( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + if IS_INTEL_ALCHEMIST and D > 128: + pytest.skip( + reason="chunk_gated_delta_rule_opt is not supported on alchemist for D>128" + ) + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + g = g / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, beta, g, h0 = map( + lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0) + ) + + tri, tri_ht = chunk_gated_delta_rule_opt( + q=( + F.normalize(q.clone(), p=2, dim=-1) + if not use_qk_l2norm_in_kernel + else q.clone() + ), + k=( + F.normalize(k.clone(), p=2, dim=-1) + if not use_qk_l2norm_in_kernel + else k.clone() + ), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + + ref, ref_ht = recurrent_gated_delta_rule_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + ) + + assert_close("o", ref, tri, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ("H", "D", "mask_p", "cu_seqlens", "dtype"), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 60, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 64, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + # bfloat16 tests + (4, 60, 0, [0, 15], torch.bfloat16), + (4, 64, 0, [0, 256, 500, 1000], torch.bfloat16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16), + ] + ], +) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", + reason="Skipping test_chunk_opt_varlen because SKIP_TEST_CHUNK_VARLEN is set", +) +def test_chunk_opt_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if IS_INTEL_ALCHEMIST and D > 128: + pytest.skip( + reason="chunk_gated_delta_rule_opt is not supported on alchemist for D>128" + ) + torch.manual_seed(42) + os.environ["TRITON_F32_DEFAULT"] = "ieee" + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, beta, g, h0 = map( + lambda x: x.to(device).requires_grad_(False), (q, k, v, beta, g, h0) + ) + + tri, tri_ht = chunk_gated_delta_rule_opt( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = recurrent_gated_delta_rule_ref( + 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]], + beta=beta[:, cu_seqlens[i] : cu_seqlens[i + 1]], + g=g[:, cu_seqlens[i] : cu_seqlens[i + 1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + assert_close("o", ref, tri, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ( + "B", + "T", + "H", + "D", + "scale", + "gate_logit_normalizer", + "mask_p", + "use_qk_l2norm_in_kernel", + "dtype", + ), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format( + *test + ), + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16), + (2, 500, 3, 60, 1, 1, 0, False, torch.bfloat16), + (2, 1000, 3, 64, 0.1, 1, 0, False, torch.bfloat16), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.bfloat16), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.bfloat16), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.bfloat16), + ] + ], +) +def test_chunk_opt_vk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + if IS_INTEL_ALCHEMIST and D > 128: + pytest.skip( + reason="chunk_gated_delta_rule_opt_vk is not supported on alchemist for D>128" + ) + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + g = g / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, beta, g, h0 = map( + lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0) + ) + + # opt_vk expects initial_state in [N, H, V, K] layout + tri, tri_ht = chunk_gated_delta_rule_opt_vk( + q=( + F.normalize(q.clone(), p=2, dim=-1) + if not use_qk_l2norm_in_kernel + else q.clone() + ), + k=( + F.normalize(k.clone(), p=2, dim=-1) + if not use_qk_l2norm_in_kernel + else k.clone() + ), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone().transpose(-1, -2).contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + + ref, ref_ht = recurrent_gated_delta_rule_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + ) + + assert_close("o", ref, tri, 0.005) + # ref_ht is [B, H, K, V], tri_ht is [N, H, V, K] + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005) + + +@pytest.mark.parametrize( + ("H", "D", "mask_p", "cu_seqlens", "dtype"), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 60, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 64, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + (4, 60, 0, [0, 15], torch.bfloat16), + (4, 64, 0, [0, 256, 500, 1000], torch.bfloat16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16), + ] + ], +) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", + reason="Skipping test_chunk_opt_vk_varlen because SKIP_TEST_CHUNK_VARLEN is set", +) +def test_chunk_opt_vk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if IS_INTEL_ALCHEMIST and D > 128: + pytest.skip( + reason="chunk_gated_delta_rule_opt_vk is not supported on alchemist for D>128" + ) + torch.manual_seed(42) + os.environ["TRITON_F32_DEFAULT"] = "ieee" + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, beta, g, h0 = map( + lambda x: x.to(device).requires_grad_(False), (q, k, v, beta, g, h0) + ) + + # opt_vk expects initial_state in [N, H, V, K] layout + tri, tri_ht = chunk_gated_delta_rule_opt_vk( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + initial_state=h0.clone().transpose(-1, -2).contiguous(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = recurrent_gated_delta_rule_ref( + 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]], + beta=beta[:, cu_seqlens[i] : cu_seqlens[i + 1]], + g=g[:, cu_seqlens[i] : cu_seqlens[i + 1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + assert_close("o", ref, tri, 0.005) + # ref_ht is [N, H, K, V], tri_ht is [N, H, V, K] + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005)