biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Misc][Model][Refactor] Pass the prefix into Linear layers (#28259)

Browse Source 
Signed-off-by: MengqingCao <cmq0113@163.com>
This commit is contained in:
Mengqing Cao
2025-11-07 19:38:38 +08:00
committed by GitHub
parent 7bdb42b2f2
commit 1958bda9b4
26 changed files with 190 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
vllm/model_executor/models/bamba.py
View File

@@ -60,6 +60,7 @@ class BambaMLP(nn.Module):
config: BambaConfig,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
prefix: str = "",
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
@@ -67,12 +68,14 @@ class BambaMLP(nn.Module):
output_sizes=[config.intermediate_size] * 2,
bias=bias,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.down_proj = RowParallelLinear(
input_size=config.intermediate_size,
output_size=config.hidden_size,
bias=bias,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
if config.hidden_act != "silu":
raise ValueError(
@@ -118,7 +121,9 @@ class BambaMixerDecoderLayer(nn.Module):
prefix=f"{prefix}.mixer",
)
self.feed_forward = BambaMLP(config, quant_config=quant_config)
self.feed_forward = BambaMLP(
config, quant_config=quant_config, prefix=f"{prefix}.feed_forward"
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -202,12 +207,14 @@ class BambaAttentionDecoderLayer(nn.Module):
self.total_num_kv_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
config.hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
)
self.attn = Attention(
@@ -219,7 +226,9 @@ class BambaAttentionDecoderLayer(nn.Module):
prefix=f"{prefix}.attn",
)
self.feed_forward = BambaMLP(config, quant_config=quant_config)
self.feed_forward = BambaMLP(
config, quant_config=quant_config, prefix=f"{prefix}.feed_forward"
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)