refactor(linear_attn): extract GDN into shared LinearAttBackend#1390
refactor(linear_attn): extract GDN into shared LinearAttBackend#1390zhangts20 wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the linear attention (GDN) mechanism for Qwen3 models by abstracting the logic into a dedicated LinearAttBackend with prefill and decode attention states. This simplifies the transformer layer inference and model-specific inference state classes. The review feedback suggests importing Qwen3NextTransformerLayerWeight from its defining module to avoid circular dependencies, and replacing .view() with .reshape() on sliced or potentially non-contiguous tensors to prevent runtime errors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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 |
There was a problem hiding this comment.
The type Qwen3NextTransformerLayerWeight is defined in lightllm.models.qwen3next.layer_weights.transformer_layer_weight, but it is imported here from transformer_layer_infer. To avoid circular dependency risks and maintain clean module boundaries, it is better to import it directly from its defining module.
| from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight | |
| from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import Qwen3NextTransformerLayerWeight |
| 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) |
There was a problem hiding this comment.
Calling .view() on a sliced tensor (like mixed_qkvzba[:, qkv_dim:z_end]) can raise a RuntimeError if the tensor is non-contiguous or if the strides are incompatible. Using .reshape() is safer and more robust as it automatically handles non-contiguous layouts by copying the data if necessary.
| z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) | |
| z = mixed_qkvzba[:, qkv_dim:z_end].reshape(-1, self.tp_num_v_heads, self.head_v_dim) |
| 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() |
There was a problem hiding this comment.
Using .view() on self.infer_state.b_req_idx can fail with a RuntimeError if the tensor is non-contiguous. It is safer to use .reshape() to ensure robustness against non-contiguous tensors.
| self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_step + 1)[:, 0].contiguous() | |
| self.b_conv_buffer_idx = self.infer_state.b_req_idx.reshape(att_batch_size, mtp_step + 1)[:, 0].contiguous() |
1 participant
No description provided.