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, 16:00 UTC
[Model] fix model testing for TeleChat2ForCausalLM and V0 llama4 (#16112)
Signed-off-by: Lu Fang <fanglu@fb.com>
Lucia Fang committed
Apr 5, 2025 at 21:23 UTC
620fc2d09ed86a4be6e9655ba94e771620a9e697
2 files changed
+11
-2
vllm/attention/backends/flash_attn.py
+5
@@ -617,10 +617,15 @@ class FlashAttentionImpl(AttentionImpl):
617
blocksparse_params: Optional[Dict[str, Any]] = None,
618
logits_soft_cap: Optional[float] = None,
619
attn_type: str = AttentionType.DECODER,
620
+ use_irope: bool = False,
621
) -> None:
622
if blocksparse_params is not None:
623
raise ValueError(
624
"FlashAttention does not support block-sparse attention.")
625
+ if use_irope:
626
+ logger.warning(
627
+ "Using irope in V0 is not supported yet, it will fall back "
628
+ "to global attention for long context.")
629
self.num_heads = num_heads
630
self.head_size = head_size
631
self.scale = float(scale)
vllm/model_executor/models/telechat2.py
+6
-2
@@ -19,7 +19,7 @@
19
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20
# See the License for the specific language governing permissions and
21
# limitations under the License.
22
-from typing import Iterable, Set, Tuple
22
+from typing import Iterable, Set, Tuple, Type
23
24
import torch
25
@@ -27,6 +27,7 @@ from vllm.config import VllmConfig
27
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
from vllm.model_executor.models.llama import LlamaForCausalLM, LlamaModel
29
30
+from .llama import LlamaDecoderLayer
31
from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
32
is_pp_missing_parameter)
33
@@ -120,7 +121,10 @@ class TeleChat2ForCausalLM(LlamaForCausalLM):
121
},
122
)
123
123
- def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
124
+ def _init_model(self,
125
+ vllm_config: VllmConfig,
126
+ prefix: str = "",
127
+ layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
128
return TeleChat2Model(vllm_config=vllm_config, prefix=prefix)
129
130
def load_weights(self, weights: Iterable[Tuple[str,