add 7b pad dim

vllm/model_executor/layers/linear.py

@@ -14,8 +14,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig, QuantizeMethodBase)
 from vllm.model_executor.utils import set_weight_attrs
-
-from vllm.logger import init_logger
 import os
 
 logger = init_logger(__name__)

vllm/model_executor/models/baichuan.py

@@ -179,7 +179,7 @@ class BaiChuanAttention(nn.Module):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.W_pack(hidden_states)
-        if os.environ.get('FA_PAD') == '1':
+        if os.environ.get('FA_PAD') == '1' and qkv.shape[-1] == 12320:
             qkv = qkv[...,:-32]
         q, k, v = qkv.chunk(chunks=3, dim=-1)
         if self.postion_embedding != "ALIBI":

vllm/model_executor/models/chatglm.py

@@ -102,7 +102,7 @@ class GLMAttention(nn.Module):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.query_key_value(hidden_states)
-        if os.environ.get('FA_PAD') == '1':
+        if os.environ.get('FA_PAD') == '1' and qkv.shape[-1] == 12320:
             qkv = qkv[...,:-32]
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self.rotary_emb(position_ids, q, k)

vllm/model_executor/models/llama.py

@@ -157,7 +157,7 @@ class LlamaAttention(nn.Module):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
-        if os.environ.get('FA_PAD') == '1':
+        if os.environ.get('FA_PAD') == '1' and qkv.shape[-1] == 12320:
             qkv = qkv[...,:-32]
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self.rotary_emb(positions, q, k)

vllm/model_executor/models/qwen2.py

@@ -149,7 +149,7 @@ class Qwen2Attention(nn.Module):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
-        if os.environ.get('FA_PAD') == '1':
+        if os.environ.get('FA_PAD') == '1' and qkv.shape[-1] == 12320:
             qkv = qkv[...,:-32]
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self.rotary_emb(positions, q, k)