diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py rename to lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py index 849bccb9e9..e972ec2844 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices, prepare_chunk_offsets -from .op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices, prepare_chunk_offsets +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp from lightllm.common.triton_utils.autotuner import autotune NUM_WARPS = [2, 4, 8, 16] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py b/lightllm/common/attention/backends/_shared/linear/cumsum.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py rename to lightllm/common/attention/backends/_shared/linear/cumsum.py index 6331e1602d..d69e44ae4c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py +++ b/lightllm/common/attention/backends/_shared/linear/cumsum.py @@ -8,12 +8,11 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .utils import check_shared_mem, input_guard +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.utils import check_shared_mem, input_guard from lightllm.common.triton_utils.autotuner import autotune BS_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/attention/backends/_shared/linear/index.py similarity index 93% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py rename to lightllm/common/attention/backends/_shared/linear/index.py index 878e2296f4..4b9ec5c001 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/attention/backends/_shared/linear/index.py @@ -8,10 +8,9 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton -from .utils import tensor_cache +from lightllm.common.attention.backends._shared.linear.utils import tensor_cache @tensor_cache diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py b/lightllm/common/attention/backends/_shared/linear/l2norm.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py rename to lightllm/common/attention/backends/_shared/linear/l2norm.py index 29f892ef26..fcad017d9b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py +++ b/lightllm/common/attention/backends/_shared/linear/l2norm.py @@ -10,9 +10,9 @@ import os import torch - import triton import triton.language as tl + from lightllm.common.triton_utils.autotuner import autotune BT_LIST = [8, 16, 32, 64, 128] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py rename to lightllm/common/attention/backends/_shared/linear/mtp_state_params.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py b/lightllm/common/attention/backends/_shared/linear/op.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py rename to lightllm/common/attention/backends/_shared/linear/op.py index 2f69aa981d..8cbde95b6b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py +++ b/lightllm/common/attention/backends/_shared/linear/op.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from .utils import is_gather_supported +from lightllm.common.attention.backends._shared.linear.utils import is_gather_supported exp = tl.exp log = tl.log diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_decode.py index f3abbeb020..03fb6e744f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from .replayssm_delta import _apply_accepted_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix def check_replay_sync_mark( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/attention/backends/_shared/linear/replayssm_delta.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_delta.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_fold.py index 42aad6edc2..2e491a038b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from .replayssm_delta import _apply_accepted_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/attention/backends/_shared/linear/solve_tril.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py rename to lightllm/common/attention/backends/_shared/linear/solve_tril.py index 3512788934..1b0a73efe0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/attention/backends/_shared/linear/solve_tril.py @@ -12,13 +12,12 @@ from typing import Optional import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import make_tensor_descriptor -from .utils import input_guard, is_amd, is_tma_supported +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import make_tensor_descriptor +from lightllm.common.attention.backends._shared.linear.utils import input_guard, is_amd, is_tma_supported def _ensure_triton_allocator(): diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py rename to lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py b/lightllm/common/attention/backends/_shared/linear/utils.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py rename to lightllm/common/attention/backends/_shared/linear/utils.py index cd7c2e3aeb..e3f6c65510 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py +++ b/lightllm/common/attention/backends/_shared/linear/utils.py @@ -16,7 +16,6 @@ from typing import Any, Literal import torch - import triton logger = logging.getLogger(__name__) diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py rename to lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py index ef38e39d6c..286b460bd1 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py +++ b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py @@ -1,11 +1,11 @@ # From : https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F from lightllm.utils.device_utils import is_tesla @@ -606,8 +606,8 @@ def torch_context_attention_fwd(q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b def test(): - import torch import numpy as np + import torch Z, H, N_CTX, D_HEAD = 16, 16, 2048, 128 dtype = torch.float16 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py index 767cee9870..002caccd4c 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py index 141587ff38..b60e4a69f4 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py @@ -1,11 +1,13 @@ +from functools import lru_cache +from typing import Dict + import torch import triton import triton.language as tl +from frozendict import frozendict + from lightllm.common.kernel_config import KernelConfigs from lightllm.utils.device_utils import calcu_kernel_best_vsm_count -from frozendict import frozendict -from functools import lru_cache -from typing import Dict class GQAVSMDecodeAttentionKernelConfig(KernelConfigs): @@ -261,7 +263,7 @@ def gqa_token_decode_attention_flash_decoding_vsm_stage1( q_head_num, kv_head_num, batch_size, - **run_config + **run_config, ): grid = (num_vsm,) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py index 5ce5d193f5..db36dc1b0c 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py @@ -1,9 +1,9 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F @triton.jit diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py index e69de29bb2..f8d0ac0a7e 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py @@ -0,0 +1,14 @@ +""" +MTP Diverse Attention Module + +MTP (Multi-Token Prediction) Diverse Attention 的实现。 +""" +from .mtp_diverse_attn import token_decode_attention_mtp_diverse_single_token +from .stage1_single_token import mtp_diverse_stage1_single_token +from .stage2_single_token import mtp_diverse_stage2_single_token + +__all__ = [ + "token_decode_attention_mtp_diverse_single_token", + "mtp_diverse_stage1_single_token", + "mtp_diverse_stage2_single_token", +] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py similarity index 83% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py index 13c7e8070e..dc98bb3544 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py @@ -1,6 +1,11 @@ import torch -from .stage1_single_token import mtp_diverse_stage1_single_token -from .stage2_single_token import mtp_diverse_stage2_single_token + +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( + mtp_diverse_stage1_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( + mtp_diverse_stage2_single_token, +) from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py index 143b3c1ce1..ba004beed3 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py @@ -12,11 +12,13 @@ - 每个请求单独做可见性检查（基于各自的 seq_len） - 中间结果按 kv block 存储，供 Stage2 聚合 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.device_utils import is_hopper from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py index 5a56426932..4d382439df 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py @@ -6,11 +6,13 @@ - 根据 seq_len 确定需要聚合的 kv block 数量 - 使用 flash attention reweighting 公式 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py index 27cb97d9a0..7db6686b8c 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py index fb122a79da..6675075c04 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py @@ -1,11 +1,12 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.kernel_config import KernelConfigs from frozendict import frozendict -from functools import lru_cache -from typing import Dict + +from lightllm.common.kernel_config import KernelConfigs class GQADiverseDecodeStage1KernelConfig(KernelConfigs): diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py index 963139b36b..54d568d578 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py @@ -1,13 +1,13 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional +from frozendict import frozendict from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict -from functools import lru_cache -from typing import Dict -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, autotune class GQADiverseDecodeStage2KernelConfig(KernelConfigs): diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py index a05db282d7..9c506d3974 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py index 9de2b82057..4065ffe6bd 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py @@ -1,8 +1,8 @@ -import torch +import math +import torch import triton import triton.language as tl -import math @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py index 96a5b26dd6..863ac322de 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py @@ -1,5 +1,4 @@ import torch - import triton import triton.language as tl diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py index d963f8582a..6248dac46d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py @@ -1,7 +1,7 @@ import torch +import torch.nn.functional as F import triton import triton.language as tl -import torch.nn.functional as F @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py index 1016e5a3c1..b6a0d37619 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py @@ -1,4 +1,5 @@ import math + import torch import triton import triton.language as tl diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/attention/backends/gdn/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_gdn_decode.py rename to lightllm/common/attention/backends/gdn/replay/decode.py index ac7af9dbdc..6dda8fc867 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/attention/backends/gdn/replay/decode.py @@ -10,11 +10,8 @@ import torch import triton +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py similarity index 91% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk.py index 7b3067bbfb..919e466472 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py @@ -10,14 +10,14 @@ import torch from einops import rearrange -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .chunk_o import chunk_fwd_o -from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from .cumsum import chunk_local_cumsum -from .l2norm import l2norm_fwd -from .solve_tril import solve_tril -from .utils import SUPPRESS_LEVEL, input_guard -from .wy_fast import recompute_w_u_fwd +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.utils import SUPPRESS_LEVEL, input_guard +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_o import chunk_fwd_o +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from lightllm.common.attention.backends.gdn.triton.kernels.wy_fast import recompute_w_u_fwd def chunk_gated_delta_rule_fwd( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py index fc49763ecd..d8c4fa91a2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py @@ -11,13 +11,12 @@ import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import exp, safe_exp -from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper from lightllm.common.triton_utils.autotuner import autotune BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py index 60a594c078..6aacf292dd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import safe_exp from lightllm.common.triton_utils.autotuner import autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py index 5dfbd6e4ab..1ee725b03d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py @@ -9,11 +9,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .op import exp +from lightllm.common.attention.backends._shared.linear.op import exp @triton.heuristics( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py b/lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py rename to lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py index 5f8cab4693..57a9d2a148 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py @@ -11,12 +11,13 @@ # https://github.com/fla-org/flash-linear-attention # ruff: noqa: E501 +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, AutotuneKernelType +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune # --------------------------------------------------------------------------- # Triton kernel diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py rename to lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py index 08bb00e644..4a8882d5fe 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py @@ -10,11 +10,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/attention/backends/kda/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_kda_decode.py rename to lightllm/common/attention/backends/kda/replay/decode.py index 3704ba4fd6..5c3d38fedc 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/attention/backends/kda/replay/decode.py @@ -10,10 +10,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/attention/backends/kda/triton/kernels/kda.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda.py index 8555b83952..54eb908af4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/attention/backends/kda/triton/kernels/kda.py @@ -24,16 +24,14 @@ import torch import triton import triton.language as tl - -from lightllm.common.triton_utils.autotuner import autotune - -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .cumsum import chunk_local_cumsum -from .index import prepare_chunk_indices -from .l2norm import l2norm_fwd from triton.language import exp2, log -from .solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.triton_utils.autotuner import autotune FLA_CHUNK_SIZE = 64 RCP_LN2 = 1.4426950216293335 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py b/lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py index be99ca9bfc..8172cd34bb 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py @@ -1,8 +1,10 @@ -from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict from functools import lru_cache from typing import Dict +from frozendict import frozendict + +from lightllm.common.kernel_config import KernelConfigs + class MlaDecodeAttentionKernelConfig(KernelConfigs): kernel_name: str = "mla_decode_attentnion" diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py index f5909fffde..63607b4594 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py @@ -1,6 +1,7 @@ import torch import triton import triton.language as tl + from lightllm.utils.device_utils import calcu_kernel_best_vsm_count diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py index 5b5e7b747c..f6f328f3e4 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py @@ -1,4 +1,5 @@ import os + import torch import triton import triton.language as tl @@ -62,7 +63,7 @@ def flash_decode_stage2( mid_out_logexpsum, B_Seqlen, Out, - **run_config + **run_config, ): if run_config: num_warps = run_config["stage2_num_warps"] diff --git a/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py index e69de29bb2..5725bed2e7 100644 --- a/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py @@ -0,0 +1 @@ +from .context_flashattention_nopad_with_v import context_attention_fwd_with_v diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py rename to lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py index be06351823..c4c539673e 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py @@ -1,9 +1,10 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F + from lightllm.utils.device_utils import is_tesla @@ -204,8 +205,8 @@ def context_attention_fwd_with_v( if __name__ == "__main__": - import torch import flashinfer + import torch Z, N_CTX, H, D_HEAD, ROPE_HEAD = 32, 1024, 16, 128, 64 dtype = torch.bfloat16 diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 137ce8ad32..2d9540645a 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -11,9 +11,9 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 2c73e97d6f..8897218fbc 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -11,9 +11,9 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..53cd88eeb1 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,23 +1,22 @@ -from lightllm.common.attention_state.bounded_att import GDNState import dataclasses import weakref -import torch from abc import ABC, abstractmethod from typing import TYPE_CHECKING -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + +import torch + +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import fused_recurrent_gated_delta_rule +from lightllm.utils.envs_utils import get_env_start_args + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -389,7 +388,7 @@ def _gdn_mtp_kernel( infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", ): - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, ) diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 11f801b62a..093ec03bef 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,6 +10,14 @@ import torch +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, @@ -17,16 +25,6 @@ BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import causal_conv1d_qkv -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 8e4a081f73..fd936ff757 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -12,34 +12,23 @@ import torch -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( - ReplayGDNState, +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, ) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.gdn.replay.decode import replay +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( - fused_gdn_gating, -) -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( - conv_pack_gdn_decode_inputs, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) from lightllm.utils.envs_utils import get_env_start_args -from .replay_gdn_decode import replay - if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index af931aa477..6275b5e4c0 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -11,30 +11,20 @@ import torch -from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, ) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.kda.replay.decode import replay +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( - chunk_kda_with_fused_gate, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) - -from .replay_kda_decode import replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py index bab3f0ad38..8da44d6ecc 100644 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ b/lightllm/common/basemodel/attention/triton/fp.py @@ -1,10 +1,13 @@ import dataclasses import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Optional + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState + class TritonAttBackend(BaseAttBackend): def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState": @@ -69,7 +72,9 @@ def _nomarl_prefill_att( att_control: AttControl = AttControl(), alloc_func=torch.empty, ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) if att_control.use_sliding_window: sliding_window = att_control.sliding_window @@ -231,9 +236,7 @@ def _spec_decode_gqa_att( v: torch.Tensor, alloc_func=torch.empty, ): - from ...triton_kernel.att.decode_att.gqa.mtp_diverse import ( - token_decode_attention_mtp_diverse_single_token, - ) + from ...triton_kernel.att.decode_att.gqa.mtp_diverse import token_decode_attention_mtp_diverse_single_token out = token_decode_attention_mtp_diverse_single_token( q=q, @@ -257,7 +260,7 @@ def _normal_decode_gqa_flash_decoding_att_vsm( alloc_func=torch.empty, ): # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_vsm import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_vsm import ( gqa_token_decode_attention_flash_decoding_vsm, ) @@ -282,7 +285,9 @@ def _normal_decode_gqa_att( alloc_func=torch.empty, ): # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.gqa_decode_flashattention_nopad import ( + gqa_decode_attention_fwd, + ) out = alloc_func(q.shape, q.dtype) @@ -312,9 +317,11 @@ def _normal_decode_stage3_att( calcu_shape1 = (batch_size, q_head_num, head_dim) att_m_tensor = alloc_func((q_head_num, total_token_num), torch.float32) - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_nopad_att1 import token_att_fwd + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( + token_attention_nopad_att1, + ) - token_att_fwd( + token_attention_nopad_att1.token_att_fwd( q.view(calcu_shape1), k, att_m_tensor, @@ -326,11 +333,11 @@ def _normal_decode_stage3_att( ) o_tensor = alloc_func(q.shape, q.dtype) - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_softmax_and_reducev import ( - token_softmax_reducev_fwd, + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( + token_attention_softmax_and_reducev, ) - token_softmax_reducev_fwd( + token_attention_softmax_and_reducev.token_softmax_reducev_fwd( att_m_tensor, v, o_tensor.view(calcu_shape1), diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py index 22e723de3b..cfb1089d10 100644 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ b/lightllm/common/basemodel/attention/triton/int4kv.py @@ -1,9 +1,12 @@ import dataclasses import weakref +from typing import Optional, Tuple + import torch + from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, Tuple + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState class Int4kvTritonAttBackend(BaseAttBackend): @@ -95,7 +98,9 @@ def _groupsize_quant_prefill_att( quant_group_size=self.backend.quant_group_size, ) - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) context_attention_fwd_contiguous_kv( q=q, diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py index 9da88626c8..b4689cb79d 100644 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ b/lightllm/common/basemodel/attention/triton/int8kv.py @@ -1,11 +1,13 @@ import dataclasses import weakref -import torch -from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Optional, Tuple -from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel + +import torch + from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers +from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel, get_env_start_args + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState class Int8kvTritonAttBackend(BaseAttBackend): @@ -94,7 +96,9 @@ def _groupsize_quant_prefill_att( quant_group_size=self.backend.quant_group_size, ) - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) context_attention_fwd_contiguous_kv( q=q, diff --git a/lightllm/common/basemodel/attention/triton/mla.py b/lightllm/common/basemodel/attention/triton/mla.py index 036e58ae70..6f5624d0f4 100644 --- a/lightllm/common/basemodel/attention/triton/mla.py +++ b/lightllm/common/basemodel/attention/triton/mla.py @@ -1,9 +1,11 @@ import dataclasses import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Tuple +import torch + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState + class MlaTritonAttBackend(BaseAttBackend): def create_att_prefill_state(self, infer_state) -> "MlaTritonPrefillAttState": @@ -77,7 +79,9 @@ def _contiguous_kv_prefill_att( att_control: AttControl, alloc_func=torch.empty, ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) k_nope, k_rope = k q_head_num = q.shape[1] diff --git a/lightllm/common/basemodel/attention/triton/neo.py b/lightllm/common/basemodel/attention/triton/neo.py index dfa31876f0..f4fb1cc030 100644 --- a/lightllm/common/basemodel/attention/triton/neo.py +++ b/lightllm/common/basemodel/attention/triton/neo.py @@ -1,10 +1,11 @@ import dataclasses import weakref -import torch from typing import Tuple +import torch + from ..base_att import AttControl -from .fp import TritonAttBackend, TritonPrefillAttState, TritonDecodeAttState +from .fp import TritonAttBackend, TritonDecodeAttState, TritonPrefillAttState class NeoTritonAttBackend(TritonAttBackend): @@ -25,7 +26,9 @@ def prefill_att( att_control: AttControl = AttControl(), alloc_func=torch.empty, ) -> torch.Tensor: - from ...triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo self.infer_state: NeoChatInferStateInfo diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index d121656c73..bac4adb5ba 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -15,8 +15,12 @@ def gqa_token_decode_attention_flash_decoding( q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) - from .gqa_flash_decoding_stage1 import flash_decode_stage1 - from .gqa_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py index f8d0ac0a7e..c759f7043d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py @@ -3,9 +3,15 @@ MTP (Multi-Token Prediction) Diverse Attention 的实现。 """ -from .mtp_diverse_attn import token_decode_attention_mtp_diverse_single_token -from .stage1_single_token import mtp_diverse_stage1_single_token -from .stage2_single_token import mtp_diverse_stage2_single_token +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.mtp_diverse_attn import ( + token_decode_attention_mtp_diverse_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( + mtp_diverse_stage1_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( + mtp_diverse_stage2_single_token, +) __all__ = [ "token_decode_attention_mtp_diverse_single_token", diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py index 936ca4dbbd..2bf2d4b86b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py @@ -1,6 +1,7 @@ -import torch from typing import Optional +import torch + def token_decode_attention_flash_decoding( q: torch.Tensor, @@ -34,7 +35,9 @@ def token_decode_attention_flash_decoding( mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=q.dtype, device="cuda") - from .int4kv_flash_decoding_stage1 import int4kv_flash_decode_stage1 + from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.int4kv_flash_decoding_stage1 import ( + int4kv_flash_decode_stage1, + ) int4kv_flash_decode_stage1( q=q.view(calcu_shape1), @@ -51,7 +54,11 @@ def token_decode_attention_flash_decoding( block_seq=BLOCK_SEQ, ) - from ..int8kv.normal.int8kv_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + int8kv_flash_decoding_stage2, + ) - flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + int8kv_flash_decoding_stage2.flash_decode_stage2( + mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ + ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py index c534795fdf..52ba55e30b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py @@ -1,9 +1,16 @@ # 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 import torch + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( + flash_decode_stage1, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + flash_diverse_decode_stage3, +) from lightllm.common.basemodel.infer_struct import InferStateInfo -from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1 -from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2 -from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3 from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py index de74f0728d..14838bc24d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py @@ -1,8 +1,14 @@ -import torch from typing import Optional + +import torch + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage1 import ( + flash_decode_stage1, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( + flash_decode_stage2, +) from lightllm.common.basemodel.infer_struct import InferStateInfo -from .int8kv_flash_decoding_stage1 import flash_decode_stage1 -from .int8kv_flash_decoding_stage2 import flash_decode_stage2 @torch.no_grad() diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py index 6c50fc3927..b61f873348 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py @@ -8,8 +8,12 @@ def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out= q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) - from .flash_decoding_stage1 import flash_decode_stage1 - from .flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py index aa7cfa9f4a..d033f7eb97 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py @@ -1,10 +1,13 @@ """Linear-attention / GDN triton kernels shared across hybrid models.""" -from .causal_conv1d import causal_conv1d_fn -from .causal_conv1d_mtp import causal_conv1d_update -from .fused_gdn_gating import fused_gdn_gating -from .gdn_decode_pack import conv_pack_gdn_decode_inputs -from .mtp_fused_recurrent import mtp_fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) + from .fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule __all__ = [ diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..16bf2b08b8 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -6,8 +6,8 @@ # The original source code was licensed under the MIT license and included # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule __all__ = [ "chunk_gated_delta_rule", diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py index 063181d995..129e943f40 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py @@ -1,12 +1,16 @@ import os +from typing import List + import torch import torch.multiprocessing as mp import triton import triton.language as tl -from typing import List -from lightllm.utils.log_utils import init_logger -from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig + +from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_config import ( + MlaDecodeAttentionKernelConfig, +) from lightllm.utils.device_utils import get_device_sm_count +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -41,8 +45,12 @@ def gqa_token_decode_attention_flash_decoding( BLOCK_N = run_config["BLOCK_N"] - from .gqa_flash_decoding_stage1 import flash_decode_stage1 - from .gqa_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q_nope.shape, q_nope.dtype, q_nope.device) if out is None else out @@ -63,7 +71,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, softmax_scale, get_sm_count=True, - **run_config + **run_config, ) if not hasattr(infer_state, "decode_att_block_seq"): @@ -111,7 +119,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, softmax_scale, get_sm_count=False, - **run_config + **run_config, ) flash_decode_stage2( @@ -121,7 +129,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), - **run_config + **run_config, ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py index 5725bed2e7..dffdd45589 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py @@ -1 +1,3 @@ -from .context_flashattention_nopad_with_v import context_attention_fwd_with_v +from lightllm.common.attention.backends.mla.triton.kernels.prefill.context_flashattention_nopad_with_v import ( + context_attention_fwd_with_v, +) diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index c60038af40..bd92c66682 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -13,8 +13,8 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.attention.linear.replay_kda_decode import get_configs, replay +from lightllm.common.attention.backends.kda.replay.decode import get_configs, replay +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, diff --git a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py index f32b093448..f1f3801522 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py @@ -1,22 +1,25 @@ -import torch -import time import os -import torch.multiprocessing as mp +import time from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage1, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( GQADiverseDecodeStage1KernelConfig, + flash_decode_stage1, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -301,9 +304,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 store_json_ans = collections.defaultdict(dict) diff --git a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py index 13c8945e59..0492884589 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py @@ -1,21 +1,24 @@ -import torch import os -import torch.multiprocessing as mp from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( GQADiverseDecodeStage2KernelConfig, + flash_decode_stage2, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -260,9 +263,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 batch_sizes = [8, 16, 32, 64] seq_lens = [32, 64, 128, 256] diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index de13eb28fc..fff9447874 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -1,16 +1,7 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - - -from typing import TYPE_CHECKING from math import prod +from typing import TYPE_CHECKING import pytest import torch @@ -18,20 +9,23 @@ from lightllm.common.attention_state.full_att import ( GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, - StaticFP8AttentionConfig, - StaticFP8FullAttentionState, PackedMLAAttentionConfig, - PackedMLALogicalConfig, PackedMLAFullAttentionState, + PackedMLALogicalConfig, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, ) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion +from unit_tests.common.att_resources.model_helpers import ( + bind_test_fp8_calibration, + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, ) if TYPE_CHECKING: from collections.abc import Sequence + from lightllm.common.attention_state.full_att.base_state import FullAttentionState @@ -335,10 +329,11 @@ def test_deferred_storage_preserves_real_fp8_calibration(kind: str): def test_int4_prefill_with_separate_storage(device, monkeypatch): import weakref from types import SimpleNamespace + + from lightllm.common.attention.backends._shared.prefill import context_flashattention_nopad from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv - from lightllm.common.basemodel.triton_kernel.att.prefill_att import context_flashattention_nopad if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py index 0989fe6d26..812fd6e9f4 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,16 +10,16 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.attention.triton import fp as state_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( mtp_diverse_attn as decode_module, stage1_single_token as stage1_module, stage2_single_token as stage2_module, ) +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.attention.triton import fp as state_module +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py index 48956c4f36..feb4a81c0d 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -7,10 +7,12 @@ import pytest import torch +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + stage2_single_token as stage2_module, +) from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py index b3068db5fb..ffc322e216 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -1,7 +1,7 @@ -import weakref import collections import inspect import math +import weakref from types import SimpleNamespace import pytest @@ -10,17 +10,19 @@ from frozendict import frozendict from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( - gqa_flash_decoding as decode_module, +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash import ( gqa_flash_decoding_stage1 as stage1_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( + gqa_flash_decoding as decode_module, +) from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py index df357539e1..5237d0b0f4 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,17 +10,17 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv import ( int4kv_flash_decoding_stage1 as stage1_module, - ppl_int4kv_flash_decoding as decode_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ppl_int4kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py index f3cb8de463..bd9619711f 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( flash_decode_stage1, ) @@ -95,7 +96,7 @@ def test_flash_decode_stage1_execution(setup_tensors): new_k = k.to(q.dtype) new_v = v.to(q.dtype) - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -174,14 +175,14 @@ def fn_triton(st=setup_tensors): del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("SUMMARY - Performance Comparison") - print(f"{'='*80}") + print(f"{'=' * 80}") print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}") - print(f"{'='*80}") + print(f"{'=' * 80}") if __name__ == "__main__": diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py index 985a33289f..449e6a83b8 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( flash_decode_stage2, ) @@ -104,7 +105,7 @@ def test_flash_decode_stage2_execution(shared_seq_len): b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"] req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :] - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -184,9 +185,9 @@ def bench_stage2(): print(results[-1]) del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"{'batch_size':<10} {'seq_len':<10} {'flash_decode_stage2_ms':<22}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<10} {r['seq_len']:<10} {r['flash_decode_stage2_ms']:<22.4f}") - print(f"{'='*80}") + print(f"{'=' * 80}") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py index c1a0ca1e58..00f0783abd 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( flash_diverse_decode_stage3, ) @@ -26,7 +27,7 @@ def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, b true_out = torch.zeros_like(out) - from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( flash_decode_stage2, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py index 7033489f66..487f03b191 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,18 +10,18 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as state_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( int8kv_flash_decoding_stage1 as stage1_module, - int8kv_flash_decoding as decode_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton import int8kv as state_module +from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import int8kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py index ecab784529..3c35d64c1a 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py @@ -23,10 +23,13 @@ """ import math + import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo +from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, +) def torch_reference_context_attention_neo( diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py index d1a53f873f..71f77aff71 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py @@ -1,7 +1,9 @@ -import torch import time + import pytest -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( +import torch + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd_contiguous_kv, ) from lightllm.utils.log_utils import init_logger diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py index 541594306d..5ddcc7c1c5 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py @@ -1,15 +1,17 @@ -import torch import time -import pytest + +import flashinfer import numpy as np +import pytest +import torch import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd, context_attention_fwd_no_prompt_cache, ) from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py index eaa19ac711..4aafd9d104 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py @@ -24,6 +24,7 @@ """ import math + import pytest import torch @@ -37,7 +38,7 @@ triton_testing = None try: - from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import ( + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( context_attention_fwd_neo, ) except ImportError: @@ -496,7 +497,7 @@ def triton_run(): f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} " f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} " f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | " - f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | " + f"fa3 {fa3_ms * 1000:8.1f} us | triton {triton_ms * 1000:8.1f} us | " f"speedup {speedup:5.2f}x" ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 956bf5e173..fac1a8bb61 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -4,9 +4,9 @@ import triton import triton.language as tl +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor def attach_replay_history( @@ -107,7 +107,7 @@ def replay_layer( """按测试 state 的类型调用对应的 replay decode。""" if state.kda: - from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + from lightllm.common.attention.backends.kda.replay.decode import replay return replay( state_ssm=state.ssm_state[layer_idx], @@ -133,7 +133,7 @@ def replay_layer( state_gate_lower_bound=state.gate_lower_bound, run_config=run_config, ) - from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + from lightllm.common.attention.backends.gdn.replay.decode import replay return replay( state=state, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 7fbb63df5d..811f66f137 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -1,9 +1,9 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, @@ -26,6 +26,7 @@ def ssm_snapshot(backend, req): from types import SimpleNamespace + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState @@ -63,7 +64,7 @@ def _replay_cu(args): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_request_counts_share_compiled_kernel(monkeypatch): - from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + from lightllm.common.attention.backends._shared.linear.ssm_acceptance import _commit_prefix, commit_ssm_prefix reqs = torch.arange(512, dtype=torch.int32, device="cuda") lengths = reqs % 4 + 1 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index f7c742db01..1aab2b32a3 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -3,7 +3,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update def causal_conv1d_ref( diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py index 969a3f354d..a67ef10cf1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -1,9 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py index 8d92d6cc8d..7e556390b5 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -9,11 +9,10 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda +from lightllm.common.attention.backends.kda.triton.kernels import kda from lightllm.common.kernel_config import KernelConfigs from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel - +from lightllm.common.triton_utils.autotuner import AutotuneLevel, Autotuner KERNELS = [ ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), @@ -212,7 +211,7 @@ def test_kda_requires_packed_inputs(fused_gate, invalid_input): def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: pytest.skip("TF32 requires SM80+") - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) lengths = (1, 16, 1519) cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) @@ -236,7 +235,7 @@ def observe(*args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") rows = torch.arange(tokens, device="cuda") % 64 matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) @@ -268,7 +267,7 @@ def observe(*args, **kwargs): ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], ) def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py index 2f8544cf75..cecbf9038c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -7,15 +7,9 @@ import torch import torch.nn.functional as F -from lightllm.common.basemodel.triton_kernel.linear_att import ( - mtp_fused_recurrent as kernel_module, -) +from lightllm.common.attention.backends.gdn.triton.kernels import mtp_fused_recurrent as kernel_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import ( - Autotuner, - AutotuneKernelType, - AutotuneLevel, -) +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner @pytest.fixture(autouse=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..bbe8580df8 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -1,10 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 8252c445a1..d03908858f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -1,11 +1,8 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig - -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 786376a362..9e9f237b75 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear.replay_gdn_decode import get_configs +from lightllm.common.attention.backends.gdn.replay.decode import get_configs from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, @@ -22,9 +22,9 @@ def ssm_snapshot(backend, req): return clone[:, req] -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 02797b61b7..a85516af28 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -1,25 +1,24 @@ +from types import SimpleNamespace + import pytest import torch +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, replay_layer, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from types import SimpleNamespace pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay + from lightllm.common.attention.backends._shared.linear.replayssm_cursor import _advance_cursor + from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} hashes = {name: set() for name in kernels} diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index a94411fbd5..bf3e383953 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning +from lightllm.common.attention.backends.gdn.replay import decode as tuning from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( attach_replay_history as ReplaySSMExecutor, ) @@ -30,7 +30,7 @@ def test_rebuild_uses_one_layer_scratch(width): @pytest.mark.parametrize("width", [1, 5]) def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): from lightllm.common.triton_utils import autotuner as autotuner_module - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index fbfc4a43a9..e7360c94ab 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -1,14 +1,13 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import KDAState -from unit_tests.common.test_linear_state import make_state -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.test_linear_state import make_state @pytest.fixture diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py index b677477519..ed33e15738 100644 --- a/unit_tests/models/glm5_next/test_kda_conv_qkv.py +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -4,15 +4,12 @@ import pytest import torch +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.l2norm import USE_DEFAULT_FLA_NORM, l2norm_fwd +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear import kda as backend_module -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( - causal_conv1d_qkv, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM - pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 49b2a08eb6..3593f5f880 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -1,29 +1,28 @@ -import weakref import dataclasses import math +import weakref from types import SimpleNamespace import pytest import torch import triton -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import set_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( +from lightllm.common.attention.backends.kda.triton.kernels.kda import ( chunk_kda_with_fused_gate, fused_kda_gate_chunk_cumsum, ) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 from lightllm.models.glm5_next.triton_kernel.kpool import ( compress_pools, - gather_pools, + expand_topk, gather_paged_pools, + gather_pools, get_pool_ranges, - expand_topk, ) -from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 - +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -412,7 +411,7 @@ def test_kpool_packed_kv_above_two_gib(): @pytest.mark.parametrize("tp_world_size", [1, 4]) def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.model import Glm5NextTpPartModel @@ -500,7 +499,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): @pytest.mark.parametrize("heads", [16, 64, 128]) def test_nope_attention_native_512_and_cuda_graph(heads): from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseDecodeState, Glm5NextSparsePrefillState if torch.cuda.get_device_capability()[0] != 9: pytest.skip("FA3 NoPE requires Hopper") diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 30c4fc22c4..6eef0d230b 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,17 +1,15 @@ import weakref + import pytest import torch +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( - fused_recurrent_kda, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -145,17 +143,11 @@ def prefill(values, destination, tail): def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.attention.linear.replay_kda import ( - ReplayKDADecodeAttState, - ) - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, - ) - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_update, - ) + from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDADecodeAttState torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index c25576a19a..32de9815e0 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -93,7 +93,7 @@ def recurrent(*, initial_state, v, **kwargs): monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + "lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp.causal_conv1d_update", conv_mtp ) monkeypatch.setattr( "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent