Chranos/enginex-mlu370-vllm
1
0
Fork 0
forked from EngineX-Cambricon/enginex-mlu370-vllm
Code Pull Requests Activity

opt llama3

Browse Source
This commit is contained in:
Chranos
2026-02-05 11:53:52 +08:00
parent 3b3e614cb6
commit 9563c9af0d
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm-v0.6.2/vllm_mlu/vllm_mlu/model_executor/models/llama.py
View File

@@ -74,18 +74,22 @@ def vllm__module_executor__models__llama__LlamaAttention__forward(
smooth_quant_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
'''
=============================
Modify by vllm_mlu
=============================
@brief: pack q & k to fit tmo.apply_rotary
@optimization: avoid redundant split operation
'''
if self.rope_scaling is not None and self.rope_scaling["rope_type"] == "longrope":
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self.rotary_emb(positions, q, k)
else:
qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
# Optimized: split qkv into [qk, v] directly, avoiding redundant split
qk, v = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
# Split qk into q and k after rotary embedding
q, k = qk.split([self.q_size, self.kv_size], dim=-1)
'''
==================
End of MLU Hijack