A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Bugfix] Revert QKVCrossParallelLinear usage in Mllama to keep BNB quantization work (#14498)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Mar 9, 2025 at 12:47 UTC
fb16eea48b4e0ac74502d1b55fb7307af2135a69
1 file changed
+29
-10
vllm/model_executor/models/mllama.py
+29
-10
@@ -43,7 +43,6 @@ from vllm.forward_context import get_forward_context
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
46
- QKVCrossParallelLinear,
46
QKVParallelLinear,
47
RowParallelLinear)
48
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -814,11 +813,20 @@ class MllamaTextCrossAttention(nn.Module):
813
self.q_local_size = self.num_local_heads * self.head_dim
814
self.kv_local_size = self.num_local_key_value_heads * self.head_dim
815
817
- self.qkv_proj = QKVCrossParallelLinear(
816
+ # TODO(Isotr0py): Use QKVCrossParallelLinear when it supports
817
+ # quantization
818
+ self.q_proj = ColumnParallelLinear(
819
+ input_size=self.hidden_size,
820
+ output_size=self.num_heads * self.head_dim,
821
+ bias=False,
822
+ quant_config=quant_config,
823
+ prefix=f"{prefix}.q_proj",
824
+ )
825
+ self.kv_proj = QKVParallelLinear(
826
self.hidden_size,
827
self.head_dim,
820
- self.num_heads,
821
- self.num_key_value_heads,
828
+ total_num_heads=0,
829
+ total_num_kv_heads=self.num_key_value_heads,
830
bias=False,
831
quant_config=quant_config,
832
prefix=f"{prefix}.qkv_proj",
@@ -854,11 +862,15 @@ class MllamaTextCrossAttention(nn.Module):
862
kv_range_for_decode: Optional[List[Tuple[int, int]]],
863
cross_attention_states: Optional[torch.Tensor],
864
) -> torch.Tensor:
857
- q, k, v = self.qkv_proj(hidden_states, cross_attention_states)
865
+ q, _ = self.q_proj(hidden_states)
866
if cross_attention_states is not None:
867
+ kv, _ = self.kv_proj(cross_attention_states)
868
+ k, v = kv.split([self.kv_local_size, self.kv_local_size], dim=-1)
869
k = k.view(-1, self.num_local_key_value_heads, self.head_dim)
870
v = v.view(-1, self.num_local_key_value_heads, self.head_dim)
871
k = self.k_norm(k)
872
+ else:
873
+ k = v = None
874
875
q = q.view(-1, self.num_local_heads, self.head_dim)
876
q = self.q_norm(q)
@@ -1149,8 +1161,13 @@ class MllamaForCausalLM(nn.Module):
1161
class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,
1162
SupportsV0Only):
1163
packed_modules_mapping = {
1152
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
1153
- "gate_up_proj": ["gate_proj", "up_proj"]
1164
+ "self_attn.qkv_proj": [
1165
+ "self_attn.q_proj",
1166
+ "self_attn.k_proj",
1167
+ "self_attn.v_proj",
1168
+ ],
1169
+ "cross_attn.kv_proj": ["cross_attn.k_proj", "cross_attn.v_proj"],
1170
+ "gate_up_proj": ["gate_proj", "up_proj"],
1171
}
1172
1173
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -1420,9 +1437,11 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,
1437
torch.Tensor]]) -> Set[str]:
1438
stacked_params_mapping = [
1439
# (param_name, shard_name, shard_id)
1423
- (".qkv_proj", ".q_proj", "q"),
1424
- (".qkv_proj", ".k_proj", "k"),
1425
- (".qkv_proj", ".v_proj", "v"),
1440
+ (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
1441
+ (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
1442
+ (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
1443
+ (".cross_attn.kv_proj", ".cross_attn.k_proj", "k"),
1444
+ (".cross_attn.kv_proj", ".cross_attn.v_proj", "v"),
1445
(".gate_up_proj", ".gate_proj", 0),
1446
(".gate_up_proj", ".up_proj", 1),
1447
]