diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 1286a46ec..55d97d2aa 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -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 diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/basemodel/attention/linear/__init__.py new file mode 100644 index 000000000..5469e1000 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/__init__.py @@ -0,0 +1,11 @@ +from .gdn import ( + LinearAttBackend, + LinearAttPrefillAttState, + LinearAttDecodeAttState, +) + +__all__ = [ + "LinearAttBackend", + "LinearAttPrefillAttState", + "LinearAttDecodeAttState", +] \ No newline at end of file diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py new file mode 100644 index 000000000..1b61c042e --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -0,0 +1,385 @@ +import dataclasses +import torch +from typing import TYPE_CHECKING +from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl +from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type +from lightllm.models.qwen3next.triton_kernel.causal_conv1d import causal_conv1d_fn +from lightllm.models.qwen3next.triton_kernel.fused_gdn_gating import fused_gdn_gating +from lightllm.models.qwen3next.triton_kernel.fla.ops import chunk_gated_delta_rule +from lightllm.models.qwen3next.triton_kernel.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.models.qwen3next.triton_kernel.mtp_fused_recurrent import mtp_fused_recurrent_gated_delta_rule +from lightllm.models.qwen3next.triton_kernel.fla.ops import fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo + from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight + + +class LinearAttBackend(BaseAttBackend): + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) + + def _init_linear_layer_metadata(self, network_config, tp_world_size): + + self.mtp_step = get_env_start_args().mtp_step + + # Linear attention specific dimensions + self.num_v_heads = network_config["linear_num_value_heads"] + self.num_k_heads = network_config["linear_num_key_heads"] + self.head_k_dim = network_config["linear_key_head_dim"] + self.head_v_dim = network_config["linear_value_head_dim"] + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] + self.activation = network_config["hidden_act"] + + # Tensor parallelism dimensions + self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size + self.tp_ba_dim = (self.num_v_heads * 2) // tp_world_size + self.tp_num_k_heads = self.num_k_heads // tp_world_size + self.tp_num_v_heads = self.num_v_heads // tp_world_size + self.tp_key_dim = self.key_dim // tp_world_size + self.tp_value_dim = self.value_dim // tp_world_size + + assert self.num_v_heads % self.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" + self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads + + # SSM state dtype optimization + ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} + start_args = get_env_start_args() + self.ssm_state_dtype = ssm_dtype_dict.get(start_args.linear_att_ssm_data_type, torch.bfloat16) + + # Pre-compute whether dtype conversion is needed + # GDN kernel output dtype is self.data_type + # Conversion needed only if SSM state uses different dtype + self.needs_ssm_dtype_conversion = get_llm_data_type() != self.ssm_state_dtype + return + + def _split_qkvzba(self, mixed_qkvzba): + qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim + z_end = qkv_dim + self.tp_value_dim + b_end = z_end + self.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return mixed_qkv, z, b, a + + def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): + if decode: + query, key, value = torch.split( + mixed_qkv, + [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], + dim=-1, + ) + batch_size = mixed_qkv.shape[0] + query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) + return query, key, value + else: + query, key, value = torch.split( + mixed_qkv, + [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], + dim=-1, + ) + seq_len = query.shape[0] + query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) + return query, key, value + + def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState": + return LinearAttPrefillAttState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": + return LinearAttDecodeAttState(backend=self, infer_state=infer_state) + + +@dataclasses.dataclass +class LinearAttPrefillAttState(BasePrefillAttState): + + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + backend: LinearAttBackend = self.backend + mtp_step = backend.mtp_step + if getattr(backend.model, "is_mtp_draft_model", False): + return + # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 + # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， + # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (mtp_step + 1) + return + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ) -> torch.Tensor: + assert att_control.linear_att_prefill, "linear_att_prefill must be True for Linear prefill attention" + assert att_control.linear_att_prefill_dict is not None, "linear_att_prefill_dict is required" + + linear_att_dict = att_control.linear_att_prefill_dict + mixed_qkvzba: torch.Tensor = linear_att_dict["mixed_qkvzba"] + layer_weight = linear_att_dict["layer_weight"] + layer_num = linear_att_dict["layer_num"] + z_out = linear_att_dict["z_out"] + + backend: LinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 + # prefill 模式下，使用不到这几个维度，所以需要扣除掉， + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) + z_out["z"] = z + + core_attn_out = self._gdn_prefill_kernel(mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight) + return core_attn_out + + def _gdn_prefill_kernel( + self, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + infer_state: "Qwen3NextInferStateInfo", + layer_weight: "Qwen3NextTransformerLayerWeight", + ): + g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) + mixed_qkv = mixed_qkv.transpose(0, 1) + + backend: LinearAttBackend = self.backend + out_tensor = causal_conv1d_fn( + mixed_qkv, + layer_weight.linear_conv1d.mm_param.weight, + bias=layer_weight.linear_conv1d.bias, + query_start_loc=infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ) + mixed_qkv = out_tensor.transpose(0, 1) + + # Recurrent processing + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + initial_state = ssm_states[self.b_ssm_buffer_idx] + # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) + core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=initial_state, + output_final_state=True, + cu_seqlens=infer_state.b1_cu_q_seq_len, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + if backend.needs_ssm_dtype_conversion: + ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(backend.ssm_state_dtype, copy=False) + else: + ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state + return core_attn_out + + +@dataclasses.dataclass +class LinearAttDecodeAttState(BaseDecodeAttState): + + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + backend: LinearAttBackend = self.backend + mtp_step = backend.mtp_step + if getattr(backend.model, "is_mtp_draft_model", False): + return + + # decode 模式下 + if mtp_step == 0: + # 非mtp模式下，不需要额外状态 + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx + return + + if mtp_step > 0: + # mtp 模式下 + batch_size = self.infer_state.batch_size + att_batch_size = batch_size // (mtp_step + 1) + assert batch_size % (mtp_step + 1) == 0 + + device = self.infer_state.b_req_idx.device + + # 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=device + ) + # shape 为 [att_batch_size] + self.b_conv_buffer_idx = self.infer_state.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=device, dtype=self.infer_state.b_req_idx.dtype + ).view(1, mtp_step + 1) + # shape 为 [att_batch_size] + # 上一步接受的数量，用于linear att 的decode mtp 算子定位正确的conv 和 ssm信息的起点。 + self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + return + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ) -> torch.Tensor: + assert att_control.linear_att_decode, "linear_att_decode must be True for Linear decode attention" + assert att_control.linear_att_decode_dict is not None, "linear_att_decode_dict is required" + + linear_att_dict = att_control.linear_att_decode_dict + mixed_qkvzba = linear_att_dict["mixed_qkvzba"] + layer_weight = linear_att_dict["layer_weight"] + layer_num = linear_att_dict["layer_num"] + z_out = linear_att_dict["z_out"] + + backend: LinearAttBackend = self.backend + + mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + + if backend.mtp_step > 0: + # MTP 模式下，使用线性层 MTP 状态。 + core_attn_out = self._gdn_mtp_kernel( + mixed_qkv, + conv_states, + ssm_states, + a, + b, + self.infer_state, + layer_weight, + ) + else: + # 非 MTP 模式下，使用线性层 decode 状态。 + core_attn_out, z = self._gdn_decode_kernel( + mixed_qkv, + z, + conv_states, + ssm_states, + a, + b, + self.infer_state, + layer_weight, + ) + z_out["z"] = z + return core_attn_out + + def _gdn_decode_kernel( + self, + mixed_qkv: torch.Tensor, + z: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + infer_state: "Qwen3NextInferStateInfo", + layer_weight: "Qwen3NextTransformerLayerWeight", + ): + backend: LinearAttBackend = self.backend + + # Recurrent processing with fused gating. Decode uses a specialized + # conv+pack kernel to avoid materializing the post-conv qkv tensor + # before immediately splitting it into q/k/v. + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + layer_weight.linear_conv1d.bias, + self.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + core_attn_out, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + inplace_final_state=True, + ssm_state_indices=self.b_ssm_buffer_idx, + use_qk_l2norm_in_kernel=True, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + a_raw=a, + b_raw=b, + ) + return core_attn_out, z + + def _gdn_mtp_kernel( + self, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + infer_state: "Qwen3NextInferStateInfo", + layer_weight: "Qwen3NextTransformerLayerWeight", + ): + from lightllm.models.qwen3next.triton_kernel.causal_conv1d_spec import ( + causal_conv1d_update as causal_conv1d_update_spec, + ) + + backend: LinearAttBackend = self.backend + + cu_seqlens_q = self.b1_mtp_cu_q_seq_len + mixed_qkv = causal_conv1d_update_spec( + mixed_qkv, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + mtp_step=backend.mtp_step, + bias=layer_weight.linear_conv1d.bias, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=cu_seqlens_q, + ) + + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) + assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" + # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, + # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step + # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q.to(torch.long), + ssm_state_indices=self.b_ssm_buffer_idx, + ssm_state_write_indices=self.b_ssm_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + a_raw=a, + b_raw=b, + ) + return core_attn_out diff --git a/lightllm/models/qwen3_5/infer_struct.py b/lightllm/models/qwen3_5/infer_struct.py index 28bded687..9bc56830e 100644 --- a/lightllm/models/qwen3_5/infer_struct.py +++ b/lightllm/models/qwen3_5/infer_struct.py @@ -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 diff --git a/lightllm/models/qwen3_5/model.py b/lightllm/models/qwen3_5/model.py index 73334be72..da1b0c1f0 100644 --- a/lightllm/models/qwen3_5/model.py +++ b/lightllm/models/qwen3_5/model.py @@ -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 diff --git a/lightllm/models/qwen3next/infer_struct.py b/lightllm/models/qwen3next/infer_struct.py index e24cd326d..5313f3c4a 100644 --- a/lightllm/models/qwen3next/infer_struct.py +++ b/lightllm/models/qwen3next/infer_struct.py @@ -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 diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index d97cc3c12..a4270ea79 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -9,17 +9,12 @@ from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.common.kv_cache_mem_manager import Qwen3NextMemManager +from lightllm.common.basemodel.attention.base_att import AttControl from typing import Tuple -from lightllm.models.qwen3next.triton_kernel.causal_conv1d import causal_conv1d_fn -from lightllm.models.qwen3next.triton_kernel.fused_gdn_gating import fused_gdn_gating -from lightllm.models.qwen3next.triton_kernel.gdn_decode_pack import conv_pack_gdn_decode_inputs from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ -from lightllm.models.qwen3next.triton_kernel.fla.ops import chunk_gated_delta_rule -from lightllm.models.qwen3next.triton_kernel.fla.ops import fused_recurrent_gated_delta_rule -from lightllm.models.qwen3next.triton_kernel.mtp_fused_recurrent import mtp_fused_recurrent_gated_delta_rule from lightllm.distributed import all_reduce from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type +from lightllm.utils.envs_utils import get_env_start_args from functools import partial logger = init_logger(__name__) @@ -44,41 +39,33 @@ def __init__(self, layer_num, network_config): num_full_attention_layers = network_config["full_attention_interval"] self.is_linear_attention_layer = (layer_num + 1) % num_full_attention_layers != 0 if self.is_linear_attention_layer: - self._init_linear_layer_metadata(layer_num, network_config) + self._init_linear_att_control() return - def _init_linear_layer_metadata(self, layer_num, network_config): - - # Linear attention specific dimensions - self.num_v_heads = network_config["linear_num_value_heads"] - self.num_k_heads = network_config["linear_num_key_heads"] - self.head_k_dim = network_config["linear_key_head_dim"] - self.head_v_dim = network_config["linear_value_head_dim"] - self.key_dim = self.head_k_dim * self.num_k_heads - self.value_dim = self.head_v_dim * self.num_v_heads - self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] - self.activation = network_config["hidden_act"] - - # Tensor parallelism dimensions - self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // self.tp_world_size_ - self.tp_ba_dim = (self.num_v_heads * 2) // self.tp_world_size_ - self.tp_num_k_heads = self.num_k_heads // self.tp_world_size_ - self.tp_num_v_heads = self.num_v_heads // self.tp_world_size_ - self.tp_key_dim = self.key_dim // self.tp_world_size_ - self.tp_value_dim = self.value_dim // self.tp_world_size_ - - assert self.num_v_heads % self.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads - - # SSM state dtype optimization - ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} - start_args = get_env_start_args() - self.ssm_state_dtype = ssm_dtype_dict.get(start_args.linear_att_ssm_data_type, torch.bfloat16) - - # Pre-compute whether dtype conversion is needed - # GDN kernel output dtype is self.data_type - # Conversion needed only if SSM state uses different dtype - self.needs_ssm_dtype_conversion = get_llm_data_type() != self.ssm_state_dtype + def _init_linear_att_control(self): + self._linear_z_out = {} + + self._linear_prefill_dict = { + "mixed_qkvzba": None, + "layer_weight": None, + "layer_num": self.layer_num_, + "z_out": self._linear_z_out, + } + self._linear_decode_dict = { + "mixed_qkvzba": None, + "layer_weight": None, + "layer_num": self.layer_num_, + "z_out": self._linear_z_out, + } + + self._linear_prefill_att_control = AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=self._linear_prefill_dict, + ) + self._linear_decode_att_control = AttControl( + linear_att_decode=True, + linear_att_decode_dict=self._linear_decode_dict, + ) return def _bind_func(self): @@ -217,7 +204,61 @@ def _get_o( o_tensor = self._tpsp_reduce(input=o_tensor, infer_state=infer_state) return o_tensor - # ==================== GDN Helper Methods ==================== + def _linear_prefill_cuda_graph_wrapper( + self, + mixed_qkvzba: torch.Tensor, + infer_state: Qwen3NextInferStateInfo, + layer_weight: Qwen3NextTransformerLayerWeight, + ) -> torch.Tensor: + + backend = infer_state.prefill_att_state1.backend + + mixed_qkvzba = mixed_qkvzba.contiguous() + _mixed_qkvzba = tensor_to_no_ref_tensor(mixed_qkvzba) + pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() + pre_capture_graph.__exit__(None, None, None) + + # _gdn_prefill_kernel returns the pre-projection value stream. Its + # logical size is num_tokens * local value heads * value head dim. + # We avoid a dry-run because FlashQLA may do host-side syncs while + # preparing varlen chunk metadata, which is illegal during capture. + num_tokens = mixed_qkvzba.shape[0] + o_shape = (num_tokens, backend.tp_num_v_heads, backend.head_v_dim) + o_dtype = mixed_qkvzba.dtype + o_device = mixed_qkvzba.device + z_shape = o_shape + + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + o = torch.empty(o_shape, dtype=o_dtype, device=o_device) + _o = tensor_to_no_ref_tensor(o) + z = torch.empty(z_shape, dtype=o_dtype, device=o_device) + _z = tensor_to_no_ref_tensor(z) + + def gdn_prefill_func(new_infer_state: "Qwen3NextInferStateInfo"): + # 与 Llama/FA3 一致：通过 new_infer_state.prefill_att_state1 进入， + # 使用 runtime 上 init_att_state() 更新过的 buffer idx，而不是 + # capture 时闭包住的 graph prefill_att_state1。 + self._linear_prefill_dict["mixed_qkvzba"] = _mixed_qkvzba + self._linear_prefill_dict["layer_weight"] = layer_weight + tmp_o = new_infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=self._linear_prefill_att_control, + alloc_func=self.alloc_tensor, + ) + # prefill_att 会把 z_out["z"] 写成新 tensor；拷回 graph buffer 后恢复引用， + # 供后续 graph 段里的 _linear_post 使用。 + _z.copy_(self._linear_z_out["z"]) + self._linear_z_out["z"] = z + tmp_o = tmp_o.view(_o.shape) + _o.copy_(tmp_o) + return + + infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=gdn_prefill_func, after_graph=pre_capture_graph) + self._linear_z_out["z"] = z + return o def context_attention_forward( self, @@ -229,7 +270,22 @@ def context_attention_forward( if not self.is_linear_attention_layer: return super().context_attention_forward(input_embdings, infer_state, layer_weight) - gdn_out = self.gdn_forward(input_embdings, infer_state, layer_weight, is_prefill=True) + assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) + mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) + + if torch.cuda.is_current_stream_capturing(): + core_attn_out = self._linear_prefill_cuda_graph_wrapper(mixed_qkvzba, infer_state, layer_weight) + else: + self._linear_prefill_dict["mixed_qkvzba"] = mixed_qkvzba + self._linear_prefill_dict["layer_weight"] = layer_weight + core_attn_out = infer_state.prefill_att_state1.prefill_att( + q=None, k=None, v=None, + att_control=self._linear_prefill_att_control, + alloc_func=self.alloc_tensor, + ) + + gdn_out = self._linear_post(core_attn_out, self._linear_z_out["z"], layer_weight) + if self.tp_world_size_ > 1: all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) return gdn_out @@ -242,54 +298,34 @@ def token_attention_forward( ): if not self.is_linear_attention_layer: return super().token_attention_forward(input_embdings, infer_state, layer_weight) - gdn_out = self.gdn_forward(input_embdings, infer_state, layer_weight, is_prefill=False) + + assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) + mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) + self._linear_decode_dict["mixed_qkvzba"] = mixed_qkvzba + self._linear_decode_dict["layer_weight"] = layer_weight + core_attn_out = infer_state.decode_att_state1.decode_att( + q=None, k=None, v=None, + att_control=self._linear_decode_att_control, + alloc_func=self.alloc_tensor, + ) + gdn_out = self._linear_post(core_attn_out, self._linear_z_out["z"], layer_weight) + if self.tp_world_size_ > 1: all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) return gdn_out - def gdn_forward( - self, - input: torch.Tensor, - infer_state: Qwen3NextInferStateInfo, - layer_weight: Qwen3NextTransformerLayerWeight, - is_prefill: bool, - ): - assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) - + def _linear_in_proj( + self, input: torch.Tensor, layer_weight: Qwen3NextTransformerLayerWeight, + ) -> torch.Tensor: input = input.view(-1, self.embed_dim_) - mixed_qkvzba = layer_weight.linear_in_proj.mm(input) - - if is_prefill: - core_attn_out, z = self._gdn_prefill_wrapper_run(mixed_qkvzba, infer_state, layer_weight) - else: - if get_env_start_args().mtp_step > 0: - # MTP 模式下，使用线性层 MTP 状态。 - mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = infer_state.req_manager.get_mamba_cache(self.layer_num_) - core_attn_out = self._gdn_mtp_kernel( - mixed_qkv, - conv_states, - ssm_states, - a, - b, - infer_state, - layer_weight, - ) - else: - # 非 MTP 模式下，使用线性层 decode 状态。 - mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = infer_state.req_manager.get_mamba_cache(self.layer_num_) - core_attn_out, z = self._gdn_decode_kernel( - mixed_qkv, - z, - conv_states, - ssm_states, - a, - b, - infer_state, - layer_weight, - ) + return layer_weight.linear_in_proj.mm(input) + def _linear_post( + self, + core_attn_out: torch.Tensor, + z: torch.Tensor, + layer_weight: Qwen3NextTransformerLayerWeight, + ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) z = z.contiguous().view(-1, z.shape[-1]) @@ -297,232 +333,3 @@ def gdn_forward( core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) return output - - def _gdn_prefill_wrapper_run( - self, - mixed_qkvzba: torch.Tensor, - infer_state: Qwen3NextInferStateInfo, - layer_weight: Qwen3NextTransformerLayerWeight, - ) -> Tuple[torch.Tensor, torch.Tensor]: - if torch.cuda.is_current_stream_capturing(): - mixed_qkvzba = mixed_qkvzba.contiguous() - _mixed_qkvzba = tensor_to_no_ref_tensor(mixed_qkvzba) - pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() - pre_capture_graph.__exit__(None, None, None) - - # _gdn_prefill_kernel returns the pre-projection value stream. Its - # logical size is num_tokens * local value heads * value head dim. - # We avoid a dry-run because FlashQLA may do host-side syncs while - # preparing varlen chunk metadata, which is illegal during capture. - num_tokens = mixed_qkvzba.shape[0] - o_shape = (num_tokens, self.tp_num_v_heads, self.head_v_dim) - o_dtype = mixed_qkvzba.dtype - o_device = mixed_qkvzba.device - z_shape = o_shape - - infer_state.prefill_cuda_graph_create_graph_obj() - infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() - o = torch.empty(o_shape, dtype=o_dtype, device=o_device) - _o = tensor_to_no_ref_tensor(o) - z = torch.empty(z_shape, dtype=o_dtype, device=o_device) - _z = tensor_to_no_ref_tensor(z) - - def gdn_prefill_func(new_infer_state: Qwen3NextInferStateInfo): - conv_states, ssm_states = new_infer_state.req_manager.get_mamba_cache(self.layer_num_) - # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 - # prefill 模式下，使用不到这几个维度，所以需要扣除掉， - if get_env_start_args().mtp_step > 0: - conv_states = conv_states[:, :, : -get_env_start_args().mtp_step] - mixed_qkv, tmp_z, b, a = self._split_qkvzba(_mixed_qkvzba) - _z.copy_(tmp_z) - tmp_o = self._gdn_prefill_kernel( - mixed_qkv, conv_states, ssm_states, a, b, new_infer_state, layer_weight - ) - tmp_o = tmp_o.view(_o.shape) - _o.copy_(tmp_o) - return - - infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=gdn_prefill_func, after_graph=pre_capture_graph) - return o, z - - conv_states, ssm_states = infer_state.req_manager.get_mamba_cache(self.layer_num_) - # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 - # prefill 模式下，使用不到这几个维度，所以需要扣除掉， - if get_env_start_args().mtp_step > 0: - conv_states = conv_states[:, :, : -get_env_start_args().mtp_step] - mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) - core_attn_out = self._gdn_prefill_kernel(mixed_qkv, conv_states, ssm_states, a, b, infer_state, layer_weight) - return core_attn_out, z - - def _split_qkvzba(self, mixed_qkvzba): - qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim - z_end = qkv_dim + self.tp_value_dim - b_end = z_end + self.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): - if decode: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - batch_size = mixed_qkv.shape[0] - query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - else: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - seq_len = query.shape[0] - query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - - def _gdn_prefill_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: Qwen3NextInferStateInfo, - layer_weight: Qwen3NextTransformerLayerWeight, - ): - g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) - mixed_qkv = mixed_qkv.transpose(0, 1) - out_tensor = causal_conv1d_fn( - mixed_qkv, - layer_weight.linear_conv1d.mm_param.weight, - bias=layer_weight.linear_conv1d.bias, - query_start_loc=infer_state.b1_cu_q_seq_len, - cache_indices=infer_state.b_conv_buffer_idx, - has_initial_state=infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation=self.activation, - ) - mixed_qkv = out_tensor.transpose(0, 1) - - # Recurrent processing - query, key, value = self._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[infer_state.b_ssm_buffer_idx] - # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( - q=query, - k=key, - v=value, - g=g.unsqueeze(0), - beta=beta.unsqueeze(0), - initial_state=initial_state, - output_final_state=True, - cu_seqlens=infer_state.b1_cu_q_seq_len, - head_first=False, - use_qk_l2norm_in_kernel=True, - ) - if self.needs_ssm_dtype_conversion: - ssm_states[infer_state.b_ssm_buffer_idx] = last_recurrent_state.to(self.ssm_state_dtype, copy=False) - else: - ssm_states[infer_state.b_ssm_buffer_idx] = last_recurrent_state - return core_attn_out - - def _gdn_decode_kernel( - self, - mixed_qkv: torch.Tensor, - z: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: Qwen3NextInferStateInfo, - layer_weight: Qwen3NextTransformerLayerWeight, - ): - # Recurrent processing with fused gating. Decode uses a specialized - # conv+pack kernel to avoid materializing the post-conv qkv tensor - # before immediately splitting it into q/k/v. - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( - mixed_qkv, - z, - a, - b, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - layer_weight.linear_conv1d.bias, - infer_state.b_conv_buffer_idx, - self.activation, - self.conv_kernel_dim, - self.tp_num_k_heads, - self.head_k_dim, - self.tp_num_v_heads, - self.head_v_dim, - ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - inplace_final_state=True, - ssm_state_indices=infer_state.b_ssm_buffer_idx, - use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out, z - - def _gdn_mtp_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: Qwen3NextInferStateInfo, - layer_weight: Qwen3NextTransformerLayerWeight, - ): - from lightllm.models.qwen3next.triton_kernel.causal_conv1d_spec import ( - causal_conv1d_update as causal_conv1d_update_spec, - ) - - cu_seqlens_q = infer_state.b1_mtp_cu_q_seq_len - mixed_qkv = causal_conv1d_update_spec( - mixed_qkv, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - mtp_step=get_env_start_args().mtp_step, - bias=layer_weight.linear_conv1d.bias, - activation=self.activation, - conv_state_indices=infer_state.b_conv_buffer_idx, - num_accepted_tokens=infer_state.b_num_accepted_tokens, - query_start_loc=cu_seqlens_q, - ) - - query, key, value = self._rearrange_mixed_qkv(mixed_qkv, decode=False) - assert infer_state.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" - # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, - # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step - # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=infer_state.b_ssm_buffer_idx, - ssm_state_write_indices=infer_state.b_ssm_buffer_idx, - num_accepted_tokens=infer_state.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 95b0b31dd..7fb795c31 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -104,3 +104,9 @@ def _init_req_manager(self): self.max_req_num, create_max_seq_len, None, linear_config=LinearAttCacheConfig.load_from_args() ) return + + def _init_att_backend1(self): + from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend + self.prefill_att_backend1 = LinearAttBackend(model=self) + self.decode_att_backend1 = LinearAttBackend(model=self) + return