[DO NOT MERGE][CAN APPROVE]change flash attn url (#1650)

deepjavalibrary · Mar 21, 2024 · cfae4c0 · cfae4c0
1 parent cd764f6
commit cfae4c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/serving/docker/deepspeed.Dockerfile b/serving/docker/deepspeed.Dockerfile
@@ -32,7 +32,7 @@ ARG deepspeed_version=nightly
 ARG deepspeed_wheel="https://publish.djl.ai/deepspeed/deepspeed-${deepspeed_version}-cp310-cp310-linux_x86_64.whl"
 # LMI-Dist Deps
 ARG vllm_wheel="https://github.com/vllm-project/vllm/releases/download/v0.3.3/vllm-0.3.3-cp310-cp310-manylinux1_x86_64.whl"
-ARG flash_attn_2_wheel="https://publish.djl.ai/flash_attn/flash_attn-2.5.6-cp310-cp310-linux_x86_64.whl"
+ARG flash_attn_2_wheel="https://publish.djl.ai/flash_attn/cu121-pt212/flash_attn-2.5.6-cp310-cp310-linux_x86_64.whl"
 ARG lmi_dist_wheel="https://publish.djl.ai/lmi_dist/lmi_dist-nightly-py3-none-any.whl"
 ARG seq_scheduler_wheel="https://publish.djl.ai/seq_scheduler/seq_scheduler-0.1.0-py3-none-any.whl"
 ARG peft_wheel="https://publish.djl.ai/peft/peft-0.5.0alpha-py3-none-any.whl"