Skip to content

Commit 60d2cfc

Browse files
youkaichaorasmith
authored andcommitted
[torch.compile] add deepseek v2 compile (vllm-project#9775)
Signed-off-by: youkaichao <youkaichao@gmail.com> Signed-off-by: Randall Smith <Randall.Smith@amd.com>
1 parent ce8ecb7 commit 60d2cfc

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

vllm/model_executor/models/deepseek_v2.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from transformers import PretrainedConfig
2929

3030
from vllm.attention import Attention, AttentionMetadata
31+
from vllm.compilation.decorators import support_torch_compile
3132
from vllm.config import CacheConfig
3233
from vllm.distributed import (get_pp_group,
3334
get_tensor_model_parallel_world_size,
@@ -403,6 +404,7 @@ def forward(
403404
return hidden_states, residual
404405

405406

407+
@support_torch_compile
406408
class DeepseekV2Model(nn.Module):
407409

408410
fall_back_to_pt_during_load = False

0 commit comments

Comments
 (0)