Skip to content

[2/5] refactor(attention): relocate kernels without changing execution #1718

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
sufubao wants to merge 1 commit into
base: refactor/attention-01-contracts
Choose a base branch
from
Draft

[2/5] refactor(attention): relocate kernels without changing execution #1718

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
0 ...triton_kernel/linear_att/causal_conv1d.py → .../backends/_shared/linear/causal_conv1d.py
View file
Open in desktop
File renamed without changes.
0 ...on_kernel/linear_att/causal_conv1d_mtp.py → ...kends/_shared/linear/causal_conv1d_mtp.py
View file
Open in desktop
File renamed without changes.
0 ...on_kernel/linear_att/causal_conv1d_qkv.py → ...kends/_shared/linear/causal_conv1d_qkv.py
View file
Open in desktop
File renamed without changes.
5 changes: 2 additions & 3 deletions ...ernel/linear_att/fla/ops/chunk_delta_h.py → .../backends/_shared/linear/chunk_delta_h.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
5 changes: 2 additions & 3 deletions ...riton_kernel/linear_att/fla/ops/cumsum.py → ...tention/backends/_shared/linear/cumsum.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
3 changes: 1 addition & 2 deletions ...triton_kernel/linear_att/fla/ops/index.py → ...ttention/backends/_shared/linear/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion ...riton_kernel/linear_att/fla/ops/l2norm.py → ...tention/backends/_shared/linear/l2norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
0 ...ton_kernel/linear_att/mtp_state_params.py → ...ckends/_shared/linear/mtp_state_params.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion ...el/triton_kernel/linear_att/fla/ops/op.py → ...n/attention/backends/_shared/linear/op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
0 ...ton_kernel/linear_att/replayssm_cursor.py → ...ckends/_shared/linear/replayssm_cursor.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion ...ton_kernel/linear_att/replayssm_decode.py → ...ckends/_shared/linear/replayssm_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
0 ...iton_kernel/linear_att/replayssm_delta.py → ...ackends/_shared/linear/replayssm_delta.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion ...riton_kernel/linear_att/replayssm_fold.py → ...backends/_shared/linear/replayssm_fold.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 3 additions & 4 deletions ...n_kernel/linear_att/fla/ops/solve_tril.py → ...ion/backends/_shared/linear/solve_tril.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
0 ...riton_kernel/linear_att/ssm_acceptance.py → ...backends/_shared/linear/ssm_acceptance.py
View file
Open in desktop
File renamed without changes.
1 change: 0 additions & 1 deletion ...triton_kernel/linear_att/fla/ops/utils.py → ...ttention/backends/_shared/linear/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from typing import Any, Literal

import torch

import triton

logger = logging.getLogger(__name__)
Expand Down
8 changes: 4 additions & 4 deletions ...efill_att/context_flashattention_nopad.py → ...d/prefill/context_flashattention_nopad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions ...ash_decoding/gqa_flash_decoding_stage1.py → ...de/gqa/flash/gqa_flash_decoding_stage1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand Down
0 ...ash_decoding/gqa_flash_decoding_stage2.py → ...de/gqa/flash/gqa_flash_decoding_stage2.py
View file
Open in desktop
File renamed without changes.
10 changes: 6 additions & 4 deletions .../flash_decoding/gqa_flash_decoding_vsm.py → ...ecode/gqa/flash/gqa_flash_decoding_vsm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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,)

Expand Down
6 changes: 3 additions & 3 deletions ...tt/gqa/gqa_decode_flashattention_nopad.py → ...de/gqa/gqa_decode_flashattention_nopad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down
14 changes: 14 additions & 0 deletions lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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",
]
9 changes: 7 additions & 2 deletions ...e_att/gqa/mtp_diverse/mtp_diverse_attn.py → ...ecode/gqa/mtp_diverse/mtp_diverse_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand Down
6 changes: 4 additions & 2 deletions ...tt/gqa/mtp_diverse/stage1_single_token.py → ...de/gqa/mtp_diverse/stage1_single_token.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
6 changes: 4 additions & 2 deletions ...tt/gqa/mtp_diverse/stage2_single_token.py → ...de/gqa/mtp_diverse/stage2_single_token.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
6 changes: 4 additions & 2 deletions ...tt/int4kv/int4kv_flash_decoding_stage1.py → ...de/int4kv/int4kv_flash_decoding_stage1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand Down
9 changes: 5 additions & 4 deletions ...v/int8kv_flash_decoding_diverse_stage1.py → ...v/int8kv_flash_decoding_diverse_stage1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down
10 changes: 5 additions & 5 deletions ...v/int8kv_flash_decoding_diverse_stage2.py → ...v/int8kv_flash_decoding_diverse_stage2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down
0 ...v/int8kv_flash_decoding_diverse_stage3.py → ...v/int8kv_flash_decoding_diverse_stage3.py
View file
Open in desktop
File renamed without changes.
6 changes: 4 additions & 2 deletions ...kv/normal/int8kv_flash_decoding_stage1.py → ...kv/normal/int8kv_flash_decoding_stage1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand Down
0 ...kv/normal/int8kv_flash_decoding_stage2.py → ...kv/normal/int8kv_flash_decoding_stage2.py
View file
Open in desktop
File renamed without changes.
0 ...a/flash_decoding/flash_decoding_stage1.py → ...decode/mha/flash/flash_decoding_stage1.py
View file
Open in desktop
File renamed without changes.
0 ...a/flash_decoding/flash_decoding_stage2.py → ...decode/mha/flash/flash_decoding_stage2.py
View file
Open in desktop
File renamed without changes.
4 changes: 2 additions & 2 deletions ..._decode_att/token_attention_nopad_att1.py → ..._decode_att/token_attention_nopad_att1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import torch
import math

import torch
import triton
import triton.language as tl
import math


@triton.jit
Expand Down
1 change: 0 additions & 1 deletion ...code_att/token_attention_nopad_reduceV.py → ...code_att/token_attention_nopad_reduceV.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import torch

import triton
import triton.language as tl

Expand Down
0 ...code_att/token_attention_nopad_softmax.py → ...code_att/token_attention_nopad_softmax.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion ...tt/token_attention_softmax_and_reducev.py → ...tt/token_attention_softmax_and_reducev.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down
1 change: 1 addition & 0 deletions .../prefill_att/context_attention_fwd_neo.py → ...nels/prefill/context_attention_fwd_neo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import math

import torch
import triton
import triton.language as tl
Expand Down
5 changes: 1 addition & 4 deletions ...del/attention/linear/replay_gdn_decode.py → ...n/attention/backends/gdn/replay/decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
16 changes: 8 additions & 8 deletions ...triton_kernel/linear_att/fla/ops/chunk.py → ...tion/backends/gdn/triton/kernels/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
7 changes: 3 additions & 4 deletions ...iton_kernel/linear_att/fla/ops/chunk_o.py → ...on/backends/gdn/triton/kernels/chunk_o.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
Loading