Skip to content

refactor(linear_attn): extract GDN into shared LinearAttBackend #1390

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

Open
zhangts20 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(linear_attn): extract GDN into shared LinearAttBackend #1390

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
5 changes: 5 additions & 0 deletions lightllm/common/basemodel/attention/base_att.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,11 @@ class AttControl:
nsa_prefill_dict: Dict = None
nsa_decode: bool = False
nsa_decode_dict: Dict = None
# linear attention 专用传参项
linear_att_prefill: bool = False
linear_att_prefill_dict: Dict = None
linear_att_decode: bool = False
linear_att_decode_dict: Dict = None


@dataclass
Expand Down
11 changes: 11 additions & 0 deletions lightllm/common/basemodel/attention/linear/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
from .gdn import (
LinearAttBackend,
LinearAttPrefillAttState,
LinearAttDecodeAttState,
)

__all__ = [
"LinearAttBackend",
"LinearAttPrefillAttState",
"LinearAttDecodeAttState",
]
385 changes: 385 additions & 0 deletions lightllm/common/basemodel/attention/linear/gdn.py
View file
Open in desktop

Large diffs are not rendered by default.

49 changes: 1 addition & 48 deletions lightllm/models/qwen3_5/infer_struct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,52 +1,5 @@
import torch

from lightllm.models.qwen2_vl.infer_struct import Qwen2VLInferStateInfo
from lightllm.utils.envs_utils import get_env_start_args


class Qwen35InferStateInfo(Qwen2VLInferStateInfo):
def __init__(self):
super().__init__()
self.gate_value = None

def init_some_extra_state(self, model):
super().init_some_extra_state(model)
mtp_step = get_env_start_args().mtp_step

is_mtp_draft_model = getattr(model, "is_mtp_draft_model", False)
if is_mtp_draft_model:
return

# prefill 模式下
if self.is_prefill:
self.b_conv_buffer_idx = self.b_req_idx
self.b_ssm_buffer_idx = self.b_req_idx * (mtp_step + 1)
return

# decode 模式下
if mtp_step == 0:
# 非mtp模式下，不需要额外状态
self.b_conv_buffer_idx = self.b_req_idx
self.b_ssm_buffer_idx = self.b_req_idx
return

if mtp_step > 0:
# mtp 模式下
batch_size = self.batch_size
att_batch_size = batch_size // (mtp_step + 1)
assert batch_size % (mtp_step + 1) == 0

# shape 为 [att_batch_size + 1]
self.b1_mtp_cu_q_seq_len = torch.arange(
0, batch_size + 1, mtp_step + 1, dtype=torch.int32, device=self.b_req_idx.device
)
# shape 为 [att_batch_size]
self.b_conv_buffer_idx = self.b_req_idx.view(att_batch_size, mtp_step + 1)[:, 0].contiguous()
self.b_ssm_buffer_idx = (self.b_conv_buffer_idx * (mtp_step + 1)).view(att_batch_size, 1) + torch.arange(
mtp_step + 1, device=self.b_req_idx.device, dtype=self.b_req_idx.dtype
).view(1, mtp_step + 1)
# shape 为 [att_batch_size]
# 上一步接受的数量，用于linear att 的decode mtp 算子定位正确的conv 和 ssm信息的起点。
self.b_num_accepted_tokens = model.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1
return
return
pass
8 changes: 8 additions & 0 deletions lightllm/models/qwen3_5/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,3 +91,11 @@ def _init_config(self):
if self.finetune_config:
self.config["vocab_size"] = self.finetune_config.vocab_size
self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)

def _init_att_backend1(self):
if getattr(self, "is_mtp_draft_model", False):
self.prefill_att_backend1 = None
self.decode_att_backend1 = None
else:
super()._init_att_backend1()
return
50 changes: 1 addition & 49 deletions lightllm/models/qwen3next/infer_struct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,53 +1,5 @@
import torch

from lightllm.models.llama.infer_struct import LlamaInferStateInfo
from lightllm.utils.envs_utils import get_env_start_args


class Qwen3NextInferStateInfo(LlamaInferStateInfo):
def __init__(self):
super().__init__()
self.gate_value = None

def init_some_extra_state(self, model):
super().init_some_extra_state(model)
mtp_step = get_env_start_args().mtp_step

is_mtp_draft_model = getattr(model, "is_mtp_draft_model", False)
if is_mtp_draft_model:
return

# prefill 模式下
if self.is_prefill:
self.b_conv_buffer_idx = self.b_req_idx
self.b_ssm_buffer_idx = self.b_req_idx * (mtp_step + 1)
return

# decode 模式下
if mtp_step == 0:
# 非mtp模式下，不需要额外状态
self.b_conv_buffer_idx = self.b_req_idx
self.b_ssm_buffer_idx = self.b_req_idx
return

if mtp_step > 0:
# mtp 模式下
batch_size = self.batch_size
att_batch_size = batch_size // (mtp_step + 1)
assert batch_size % (mtp_step + 1) == 0

# shape 为 [att_batch_size + 1]
self.b1_mtp_cu_q_seq_len = torch.arange(
0, batch_size + 1, mtp_step + 1, dtype=torch.int32, device=self.b_req_idx.device
)
# shape 为 [att_batch_size]
self.b_conv_buffer_idx = self.b_req_idx.view(att_batch_size, mtp_step + 1)[:, 0].contiguous()
# shape 为 [att_batch_size, mtp_step + 1]
self.b_ssm_buffer_idx = (self.b_conv_buffer_idx * (mtp_step + 1)).view(att_batch_size, 1) + torch.arange(
mtp_step + 1, device=self.b_req_idx.device, dtype=self.b_req_idx.dtype
).view(1, mtp_step + 1)
# shape 为 [att_batch_size]
# 上一步接受的数量，用于linear att 的decode mtp 算子定位正确的conv 和 ssm信息的起点。
self.b_num_accepted_tokens = model.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1
return
return
pass
Loading
Loading