Unverified Commit 77e6dcbb authored by Shanshan Shen's avatar Shanshan Shen Committed by GitHub
Browse files

[PluggableLayer][MM] Add PluggableLayer for RelPosAttention (#33753)


Signed-off-by: default avatarshen-shanshan <467638484@qq.com>
parent 70c73df6
...@@ -54,6 +54,8 @@ For example: ...@@ -54,6 +54,8 @@ For example:
--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn" --8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
--8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention" --8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
--8<-- "vllm/model_executor/models/deepencoder.py:rel_pos_attention"
``` ```
**2. Activation:** **2. Activation:**
......
...@@ -18,6 +18,7 @@ import torch.nn as nn ...@@ -18,6 +18,7 @@ import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from transformers import CLIPVisionConfig from transformers import CLIPVisionConfig
from vllm.model_executor.custom_op import PluggableLayer
from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.attention import MMEncoderAttention
from vllm.model_executor.layers.conv import Conv2dLayer from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization import QuantizationConfig
...@@ -263,9 +264,13 @@ class Block(nn.Module): ...@@ -263,9 +264,13 @@ class Block(nn.Module):
return x return x
class RelPosAttention(nn.Module): # --8<-- [start:rel_pos_attention]
@PluggableLayer.register("rel_pos_attention")
class RelPosAttention(PluggableLayer):
"""Multi-head Attention block with relative position embeddings.""" """Multi-head Attention block with relative position embeddings."""
# --8<-- [end:rel_pos_attention]
def __init__( def __init__(
self, self,
dim: int, dim: int,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment