From 19800d1cadd79e87d5cb20ee36736b81d1ccffed Mon Sep 17 00:00:00 2001 From: Satya Nikhil Kodukula Date: Tue, 29 Sep 2026 12:49:58 +0000 Subject: [PATCH] Move _triton_kernels/gated_delta_rule/ to _triton_kernels/gated_delta_net/ The wrappers that launch these kernels live in aiter/ops/triton/gated_delta_net/, so the kernel directory takes the same name; kernel directories mirror the wrapper folder, not the kernel family. git mv of the 26 files plus the import path in every module that names the old one: the gated_delta_net wrappers, the kernels' own absolute imports, chunk_delta_attn/chunk_fwd.py, aiter/ops/chunk_gated_delta_rule_fwd_h.py, aiter/ops/flydsl, the two benchmarks and the three op_tests scripts. No code change otherwise. --- aiter/ops/chunk_gated_delta_rule_fwd_h.py | 4 ++-- .../flydsl/linear_attention_prefill_kernels.py | 2 +- .../chunk_delta_attn/chunk_fwd.py | 4 ++-- .../__init__.py | 14 +++++++------- .../decode/__init__.py | 6 +++--- .../decode/causal_conv1d_update_split_qkv.py | 0 .../decode/fused_conv_recurrent_norm.py | 0 .../decode/fused_conv_recurrent_norm_unified.py | 0 .../decode/fused_rearrange_sigmoid_gdr.py | 0 .../decode/fused_recurrent.py | 0 .../decode/fused_sigmoid_gating_recurrent.py | 2 +- .../fused_qkvzba_split.py | 0 .../gated_delta_rule_utils.py | 0 .../prefill/__init__.py | 12 ++++++------ .../prefill/causal_conv1d_fwd_split_qkv.py | 0 .../prefill/chunk.py | 10 +++++----- .../prefill/chunk_delta_h.py | 6 +++--- .../prefill/chunk_o.py | 6 +++--- .../prefill/fused_cumsum_kkt.py | 6 +++--- .../prefill/fused_gdn_gating_prefill.py | 0 .../prefill/fused_solve_tril_recompute.py | 8 ++++---- .../utils/__init__.py | 12 ++++++------ .../utils/cumsum.py | 4 ++-- .../utils/index.py | 2 +- .../utils/l2norm.py | 2 +- .../utils/op.py | 2 +- .../utils/prefill_metadata.py | 0 .../utils/solve_tril.py | 6 +++--- .../utils/wy_representation.py | 6 +++--- aiter/ops/triton/gated_delta_net/__init__.py | 2 +- .../gated_delta_net/causal_conv1d_decode.py | 4 ++-- .../gated_delta_net/causal_conv1d_prefill.py | 4 ++-- .../triton/gated_delta_net/fused_kda_decode.py | 2 +- .../gated_delta_net/fused_kda_decode_unified.py | 2 +- .../fused_rearrange_sigmoid_gdr.py | 2 +- .../triton/gated_delta_net/gated_delta_rule.py | 4 ++-- .../triton/bench_fused_kda_decode.py | 2 +- .../bench_gated_delta_rule_snapshot_dtype.py | 2 +- op_tests/test_flydsl_linear_attention_prefill.py | 2 +- op_tests/test_gated_delta_rule.py | 16 ++++++++-------- op_tests/test_gdn_prepare.py | 4 ++-- 41 files changed, 80 insertions(+), 80 deletions(-) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/__init__.py (73%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/__init__.py (69%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/causal_conv1d_update_split_qkv.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/fused_conv_recurrent_norm.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/fused_conv_recurrent_norm_unified.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/fused_rearrange_sigmoid_gdr.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/fused_recurrent.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/decode/fused_sigmoid_gating_recurrent.py (98%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/fused_qkvzba_split.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/gated_delta_rule_utils.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/__init__.py (70%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/causal_conv1d_fwd_split_qkv.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/chunk.py (97%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/chunk_delta_h.py (99%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/chunk_o.py (99%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/fused_cumsum_kkt.py (97%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/fused_gdn_gating_prefill.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/prefill/fused_solve_tril_recompute.py (98%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/__init__.py (67%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/cumsum.py (98%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/index.py (98%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/l2norm.py (99%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/op.py (95%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/prefill_metadata.py (100%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/solve_tril.py (99%) rename aiter/ops/triton/_triton_kernels/{gated_delta_rule => gated_delta_net}/utils/wy_representation.py (97%) diff --git a/aiter/ops/chunk_gated_delta_rule_fwd_h.py b/aiter/ops/chunk_gated_delta_rule_fwd_h.py index b59e5fe6d19..12447e91ec8 100644 --- a/aiter/ops/chunk_gated_delta_rule_fwd_h.py +++ b/aiter/ops/chunk_gated_delta_rule_fwd_h.py @@ -9,7 +9,7 @@ from torch import Tensor from ..jit.core import compile_ops -from .triton._triton_kernels.gated_delta_rule.utils.prefill_metadata import ( +from .triton._triton_kernels.gated_delta_net.utils.prefill_metadata import ( GatedDeltaRulePrefillMetadata, build_gated_delta_rule_prefill_metadata, ) @@ -297,7 +297,7 @@ def chunk_gated_delta_rule_fwd_h_hip_fn( ) if is_varlen: - from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( + from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( prepare_chunk_offsets, prepare_rebased_cu_seqlens, ) diff --git a/aiter/ops/flydsl/linear_attention_prefill_kernels.py b/aiter/ops/flydsl/linear_attention_prefill_kernels.py index 4bfbb64f050..f7029b3ba9e 100644 --- a/aiter/ops/flydsl/linear_attention_prefill_kernels.py +++ b/aiter/ops/flydsl/linear_attention_prefill_kernels.py @@ -28,7 +28,7 @@ from aiter.jit.core import AITER_CONFIGS -from ..triton._triton_kernels.gated_delta_rule.utils import ( +from ..triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, build_gated_delta_rule_prefill_metadata, prepare_chunk_offsets, diff --git a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py index 5431138d52c..c5397e81aa4 100644 --- a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py +++ b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py @@ -38,7 +38,7 @@ prepare_chunk_indices, ) from aiter.ops.triton._triton_kernels.chunk_delta_attn.utils.l2norm import l2norm_fwd -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk_delta_h import ( chunk_gated_delta_rule_fwd_h, ) @@ -218,7 +218,7 @@ def chunk_delta_attn_fwd( lower_bound=lower_bound, ) else: - from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( + from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( chunk_local_cumsum, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/__init__.py similarity index 73% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/__init__.py index 523c56230c7..0790ed82b8f 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/__init__.py @@ -17,29 +17,29 @@ For training with gradients, please use the flash-linear-attention library. """ -from aiter.ops.triton._triton_kernels.gated_delta_rule import gated_delta_rule_utils -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_recurrent import ( +from aiter.ops.triton._triton_kernels.gated_delta_net import gated_delta_rule_utils +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_recurrent import ( _fused_recurrent_gated_delta_rule_fwd_kernel, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk import ( 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.prefill.chunk_delta_h import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.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 aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_o import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk_o import ( chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_cumsum_kkt import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_cumsum_kkt import ( fused_chunk_local_cumsum_scaled_dot_kkt_fwd, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_solve_tril_recompute import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_solve_tril_recompute import ( fused_solve_tril_recompute_w_u, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/__init__.py similarity index 69% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/__init__.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/__init__.py index 4e69db8b637..6a5e2df6e03 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/__init__.py @@ -8,13 +8,13 @@ This module provides optimized Triton kernels for decode/inference operations. """ -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_rearrange_sigmoid_gdr import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_rearrange_sigmoid_gdr import ( fused_rearrange_sigmoid_gated_delta_rule_update_kernel, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_recurrent import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_recurrent import ( _fused_recurrent_gated_delta_rule_fwd_kernel, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/causal_conv1d_update_split_qkv.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/causal_conv1d_update_split_qkv.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/causal_conv1d_update_split_qkv.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/causal_conv1d_update_split_qkv.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_conv_recurrent_norm.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_conv_recurrent_norm.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_conv_recurrent_norm.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_conv_recurrent_norm.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_conv_recurrent_norm_unified.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_conv_recurrent_norm_unified.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_conv_recurrent_norm_unified.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_conv_recurrent_norm_unified.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_rearrange_sigmoid_gdr.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_rearrange_sigmoid_gdr.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_rearrange_sigmoid_gdr.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_rearrange_sigmoid_gdr.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_recurrent.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_recurrent.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_recurrent.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_recurrent.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_sigmoid_gating_recurrent.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_sigmoid_gating_recurrent.py similarity index 98% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_sigmoid_gating_recurrent.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_sigmoid_gating_recurrent.py index e02e8d037da..75f3b7d3c42 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/decode/fused_sigmoid_gating_recurrent.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/decode/fused_sigmoid_gating_recurrent.py @@ -2,7 +2,7 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( autotune_cache_kwargs, input_guard, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/fused_qkvzba_split.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/fused_qkvzba_split.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/fused_qkvzba_split.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/fused_qkvzba_split.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/gated_delta_rule_utils.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/gated_delta_rule_utils.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/gated_delta_rule_utils.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/__init__.py similarity index 70% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/__init__.py index fa4fa271a2c..6732e2e6e12 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/__init__.py @@ -8,29 +8,29 @@ This module provides optimized Triton kernels for prefill/training operations. """ -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk import ( 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.prefill.chunk_delta_h import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.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 aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_o import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk_o import ( chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_cumsum_kkt import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_cumsum_kkt import ( fused_chunk_local_cumsum_scaled_dot_kkt_fwd, fused_cumsum_kkt, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_gdn_gating_prefill import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_gdn_gating_prefill import ( fused_gdn_gating_and_sigmoid, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_solve_tril_recompute import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_solve_tril_recompute import ( fused_solve_tril_recompute_w_u, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/causal_conv1d_fwd_split_qkv.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/causal_conv1d_fwd_split_qkv.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/causal_conv1d_fwd_split_qkv.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/causal_conv1d_fwd_split_qkv.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk.py similarity index 97% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk.py index a52962f2cfc..ad2a23c34a2 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk.py @@ -14,23 +14,23 @@ import torch -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.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 aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_o import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk_o import ( chunk_fwd_o, chunk_fwd_o_opt, chunk_fwd_o_opt_vk, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_cumsum_kkt import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_cumsum_kkt import ( fused_chunk_local_cumsum_scaled_dot_kkt_fwd, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_solve_tril_recompute import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.fused_solve_tril_recompute import ( fused_solve_tril_recompute_w_u, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, build_gated_delta_rule_prefill_metadata, chunk_local_cumsum, diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_delta_h.py similarity index 99% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_delta_h.py index 42728a6ed2c..2602d10c598 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_delta_h.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_delta_h.py @@ -14,20 +14,20 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_AMD, IS_NVIDIA_HOPPER, RCP_LN2, autotune_cache_kwargs, check_shared_mem, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, prepare_chunk_indices, prepare_chunk_offsets, prepare_rebased_cu_seqlens, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import exp +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import exp from aiter.ops.triton.utils.tuned_config_utils import autotune_configs NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_o.py similarity index 99% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_o.py index 703e9159e35..1bdf4accbcf 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/chunk_o.py @@ -12,17 +12,17 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, prepare_chunk_indices, prepare_rebased_cu_seqlens, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import exp +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import exp from aiter.ops.triton.utils.tuned_config_utils import autotune_configs BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_cumsum_kkt.py similarity index 97% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_cumsum_kkt.py index 8b36299d93b..a4aadd1c7e3 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_cumsum_kkt.py @@ -2,17 +2,17 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_AMD, RCP_LN2, autotune_cache_kwargs, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, prepare_chunk_indices, prepare_rebased_cu_seqlens, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import exp +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import exp from aiter.ops.triton.utils.tuned_config_utils import autotune_configs diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_gdn_gating_prefill.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_gdn_gating_prefill.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_gdn_gating_prefill.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_gdn_gating_prefill.py 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_net/prefill/fused_solve_tril_recompute.py similarity index 98% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_solve_tril_recompute.py index 19f4959ed80..4bab938b744 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/prefill/fused_solve_tril_recompute.py @@ -15,17 +15,17 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_AMD, autotune_cache_kwargs, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, prepare_chunk_indices, prepare_rebased_cu_seqlens, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import exp -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.solve_tril import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import exp +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.solve_tril import ( FLA_TRIL_PRECISION, solve_tril, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/__init__.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/__init__.py similarity index 67% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/__init__.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/__init__.py index 10ce91a3ad3..edea73c304b 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/__init__.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/__init__.py @@ -2,29 +2,29 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. # Adapted from flash-linear-attention: Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.cumsum import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.cumsum import ( chunk_local_cumsum, chunk_local_cumsum_scalar, chunk_local_cumsum_vector, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.index import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.index import ( prepare_chunk_indices, prepare_chunk_offsets, prepare_num_chunks, prepare_rebased_cu_seqlens, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.l2norm import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.l2norm import ( l2norm_bwd, l2norm_fwd, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.prefill_metadata import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.prefill_metadata import ( GatedDeltaRulePrefillMetadata, build_gated_delta_rule_prefill_metadata, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.solve_tril import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.solve_tril import ( solve_tril, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.wy_representation import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.wy_representation import ( chunk_scaled_dot_kkt_fwd, recompute_w_u_fwd, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/cumsum.py similarity index 98% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/cumsum.py index f557b112af4..4bbb0e61198 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/cumsum.py @@ -13,12 +13,12 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( autotune_cache_kwargs, check_shared_mem, input_guard, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.index import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.index import ( prepare_chunk_indices, ) from aiter.ops.triton.utils.tuned_config_utils import autotune_configs diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/index.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/index.py similarity index 98% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/index.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/index.py index fb94981c602..2ac5ef42a65 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/index.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/index.py @@ -13,7 +13,7 @@ import torch.nn.functional as F import triton -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( tensor_cache, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/l2norm.py similarity index 99% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/l2norm.py index c2c25984b69..3b5a6c9d896 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/l2norm.py @@ -14,7 +14,7 @@ import triton.language as tl from torch import nn -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_AMD, autotune_cache_kwargs, input_guard, diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/op.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/op.py similarity index 95% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/op.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/op.py index 6d23ae92b72..0898843a5d1 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/op.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/op.py @@ -15,7 +15,7 @@ import triton.language as tl import triton.language.extra.libdevice as tldevice -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_GATHER_SUPPORTED, IS_TMA_SUPPORTED, ) diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/prefill_metadata.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/prefill_metadata.py similarity index 100% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/prefill_metadata.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/prefill_metadata.py diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/solve_tril.py similarity index 99% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/solve_tril.py index 61708242bc8..53d7f811514 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/solve_tril.py @@ -15,15 +15,15 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.index import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.index import ( prepare_chunk_indices, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import ( make_tensor_descriptor, ) from aiter.ops.triton.utils.tuned_config_utils import autotune_configs diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/wy_representation.py similarity index 97% rename from aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py rename to aiter/ops/triton/_triton_kernels/gated_delta_net/utils/wy_representation.py index d9c1cfab573..1d3e2178baa 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_net/utils/wy_representation.py @@ -13,13 +13,13 @@ import triton import triton.language as tl -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( autotune_cache_kwargs, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.index import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.index import ( prepare_chunk_indices, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils.op import exp +from aiter.ops.triton._triton_kernels.gated_delta_net.utils.op import exp from aiter.ops.triton.utils.tuned_config_utils import autotune_configs diff --git a/aiter/ops/triton/gated_delta_net/__init__.py b/aiter/ops/triton/gated_delta_net/__init__.py index d48c8ca3a8c..8f0a375069b 100644 --- a/aiter/ops/triton/gated_delta_net/__init__.py +++ b/aiter/ops/triton/gated_delta_net/__init__.py @@ -8,7 +8,7 @@ This module provides high-level Triton implementations for gated delta rule. """ -from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, build_gated_delta_rule_prefill_metadata, ) diff --git a/aiter/ops/triton/gated_delta_net/causal_conv1d_decode.py b/aiter/ops/triton/gated_delta_net/causal_conv1d_decode.py index 09ced58e3a0..804cfca0a9c 100644 --- a/aiter/ops/triton/gated_delta_net/causal_conv1d_decode.py +++ b/aiter/ops/triton/gated_delta_net/causal_conv1d_decode.py @@ -6,13 +6,13 @@ Decode-stage counterpart of the prefill ``causal_conv1d_prefill`` module: a single-/few-token autoregressive conv update that directly outputs split q/k/v. The ``@triton.jit`` / Gluon kernels live in -``aiter.ops.triton._triton_kernels.gated_delta_rule.decode.causal_conv1d_update_split_qkv``. +``aiter.ops.triton._triton_kernels.gated_delta_net.decode.causal_conv1d_update_split_qkv``. """ import torch import triton -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.causal_conv1d_update_split_qkv import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.causal_conv1d_update_split_qkv import ( PAD_SLOT_ID, _causal_conv1d_update_split_qkv_kernel, gluon_causal_conv1d_update_split_qkv_kernel, diff --git a/aiter/ops/triton/gated_delta_net/causal_conv1d_prefill.py b/aiter/ops/triton/gated_delta_net/causal_conv1d_prefill.py index bb2e0421300..2b03a3fd451 100644 --- a/aiter/ops/triton/gated_delta_net/causal_conv1d_prefill.py +++ b/aiter/ops/triton/gated_delta_net/causal_conv1d_prefill.py @@ -17,7 +17,7 @@ * FlyDSL : ``aiter.ops.flydsl.causal_conv1d_flydsl.causal_conv1d_split_qkv_flydsl_fn`` The ``@triton.jit`` kernels themselves live in -``aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.causal_conv1d_fwd_split_qkv``. +``aiter.ops.triton._triton_kernels.gated_delta_net.prefill.causal_conv1d_fwd_split_qkv``. ``x`` has shape ``[dim, cu_seqlen]`` with the channel axis packed as the concatenation ``[Q | K | V]`` (``dim == 2*k_dim + v_dim``). @@ -27,7 +27,7 @@ import triton from aiter.ops.prefill_batch_metadata import CausalConvPrefillMetadata -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.causal_conv1d_fwd_split_qkv import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.causal_conv1d_fwd_split_qkv import ( PAD_SLOT_ID, _causal_conv1d_fwd_split_qkv_kernel, _causal_conv1d_fwd_split_qkv_tile_kernel, diff --git a/aiter/ops/triton/gated_delta_net/fused_kda_decode.py b/aiter/ops/triton/gated_delta_net/fused_kda_decode.py index ed0996bbbce..2cd9d81c003 100644 --- a/aiter/ops/triton/gated_delta_net/fused_kda_decode.py +++ b/aiter/ops/triton/gated_delta_net/fused_kda_decode.py @@ -9,7 +9,7 @@ import torch -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_conv_recurrent_norm import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_conv_recurrent_norm import ( fused_conv_recurrent_norm_kernel, ) from aiter.ops.triton.utils._triton.arch_info import get_arch diff --git a/aiter/ops/triton/gated_delta_net/fused_kda_decode_unified.py b/aiter/ops/triton/gated_delta_net/fused_kda_decode_unified.py index f646df27996..f30cb2742da 100644 --- a/aiter/ops/triton/gated_delta_net/fused_kda_decode_unified.py +++ b/aiter/ops/triton/gated_delta_net/fused_kda_decode_unified.py @@ -15,7 +15,7 @@ import torch import triton -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_conv_recurrent_norm_unified import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_conv_recurrent_norm_unified import ( _fused_kda_decode_unified_kernel, ) from aiter.ops.triton.utils._triton.arch_info import get_arch diff --git a/aiter/ops/triton/gated_delta_net/fused_rearrange_sigmoid_gdr.py b/aiter/ops/triton/gated_delta_net/fused_rearrange_sigmoid_gdr.py index 770e9c2a4c4..64fd510f169 100644 --- a/aiter/ops/triton/gated_delta_net/fused_rearrange_sigmoid_gdr.py +++ b/aiter/ops/triton/gated_delta_net/fused_rearrange_sigmoid_gdr.py @@ -12,7 +12,7 @@ import torch import triton -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_rearrange_sigmoid_gdr import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_rearrange_sigmoid_gdr import ( fused_rearrange_sigmoid_gated_delta_rule_update_kernel, ) 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 9b18cae5054..b8815b52160 100644 --- a/aiter/ops/triton/gated_delta_net/gated_delta_rule.py +++ b/aiter/ops/triton/gated_delta_net/gated_delta_rule.py @@ -21,13 +21,13 @@ import torch import triton -from aiter.ops.triton._triton_kernels.gated_delta_rule import ( +from aiter.ops.triton._triton_kernels.gated_delta_net 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 ( +from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( GatedDeltaRulePrefillMetadata, l2norm_fwd, ) diff --git a/op_tests/op_benchmarks/triton/bench_fused_kda_decode.py b/op_tests/op_benchmarks/triton/bench_fused_kda_decode.py index c2a34360f90..25ef45ce456 100644 --- a/op_tests/op_benchmarks/triton/bench_fused_kda_decode.py +++ b/op_tests/op_benchmarks/triton/bench_fused_kda_decode.py @@ -21,7 +21,7 @@ import triton.language as tl from einops import rearrange -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) from aiter.ops.triton.gated_delta_net.causal_conv1d_decode import ( diff --git a/op_tests/op_benchmarks/triton/bench_gated_delta_rule_snapshot_dtype.py b/op_tests/op_benchmarks/triton/bench_gated_delta_rule_snapshot_dtype.py index 70e55820e32..59400e8db19 100644 --- a/op_tests/op_benchmarks/triton/bench_gated_delta_rule_snapshot_dtype.py +++ b/op_tests/op_benchmarks/triton/bench_gated_delta_rule_snapshot_dtype.py @@ -20,7 +20,7 @@ from aiter.ops.chunk_gated_delta_rule_fwd_h import ( chunk_gated_delta_rule_fwd_h_hip, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill import ( chunk_fwd_o_opt_vk, ) diff --git a/op_tests/test_flydsl_linear_attention_prefill.py b/op_tests/test_flydsl_linear_attention_prefill.py index e587b17e4b8..17fe3b51d4e 100644 --- a/op_tests/test_flydsl_linear_attention_prefill.py +++ b/op_tests/test_flydsl_linear_attention_prefill.py @@ -47,7 +47,7 @@ from aiter.ops.prefill_batch_metadata import ( build_gated_delta_rule_prefill_metadata, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk_delta_h import ( chunk_gated_delta_rule_fwd_h_opt_vk, ) from aiter.test_common import benchmark, checkAllclose, run_perftest diff --git a/op_tests/test_gated_delta_rule.py b/op_tests/test_gated_delta_rule.py index 42867b4d543..453e52af819 100644 --- a/op_tests/test_gated_delta_rule.py +++ b/op_tests/test_gated_delta_rule.py @@ -21,21 +21,21 @@ from aiter.ops.flydsl.linear_attention_prefill_kernels import ( chunk_gated_delta_rule_fwd_h_flydsl_opt, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.decode.fused_sigmoid_gating_recurrent import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.decode.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.gated_delta_rule_utils import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.gated_delta_rule_utils import ( IS_AMD, IS_INTEL_ALCHEMIST, assert_close, device, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill import ( chunk_gated_delta_rule_fwd_h_opt_vk, fused_chunk_local_cumsum_scaled_dot_kkt_fwd, fused_solve_tril_recompute_w_u, ) -from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( +from aiter.ops.triton._triton_kernels.gated_delta_net.prefill import ( fused_solve_tril_recompute as fused_solve_module, ) from aiter.ops.triton.gated_delta_net import ( @@ -94,7 +94,7 @@ def test_chunk_opt_vk_unsupported_gfx12_runtime_allowlist( monkeypatch, arch: str, expected: bool ): chunk_module = importlib.import_module( - "aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk" + "aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk" ) props = type("DeviceProperties", (), {"gcnArchName": arch})() monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: props) @@ -1627,7 +1627,7 @@ def test_chunk_fwd_h_beyond_int32_chunk_offsets(seqlens): faulted the GPU. Zero k/w/u and no gate leave the state untouched, so every chunk snapshot must come back exactly equal to the initial state. """ - from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( + from aiter.ops.triton._triton_kernels.gated_delta_net.prefill import ( chunk_gated_delta_rule_fwd_h, ) @@ -1835,7 +1835,7 @@ def test_chunk_opt_vk_rejects_dense_index_count_mismatch(): def test_chunk_opt_vk_hip_downgrade_preserves_indexed_state_args(monkeypatch): chunk_module = importlib.import_module( - "aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk" + "aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk" ) initial_state = torch.empty(4, 1, 1, 1) initial_state_indices = torch.tensor([3], dtype=torch.int32) @@ -1893,7 +1893,7 @@ def fake_triton_k5(**kwargs): def test_chunk_opt_vk_unsupported_gfx12_downgrades_to_triton(monkeypatch): chunk_module = importlib.import_module( - "aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk" + "aiter.ops.triton._triton_kernels.gated_delta_net.prefill.chunk" ) props = type("DeviceProperties", (), {"gcnArchName": "gfx1250"})() monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: props) diff --git a/op_tests/test_gdn_prepare.py b/op_tests/test_gdn_prepare.py index 480bfc6df3a..6fe65436a68 100644 --- a/op_tests/test_gdn_prepare.py +++ b/op_tests/test_gdn_prepare.py @@ -101,11 +101,11 @@ def _load_kernels(): gdn_prepare_flydsl_supported, gdn_prepare_fwd_flydsl, ) - from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill import ( + from aiter.ops.triton._triton_kernels.gated_delta_net.prefill import ( fused_chunk_local_cumsum_scaled_dot_kkt_fwd, fused_solve_tril_recompute_w_u, ) - from aiter.ops.triton._triton_kernels.gated_delta_rule.utils import ( + from aiter.ops.triton._triton_kernels.gated_delta_net.utils import ( build_gated_delta_rule_prefill_metadata, ) from aiter.ops.triton.gated_delta_net import chunk_gated_delta_rule_opt_vk