Re-apply vllm-project#17446

Signed-off-by: Gregory Shtrasberg <[email protected]>

Author: gshtras

vllm/attention/ops/triton_flash_attention.py

@@ -22,11 +22,10 @@
 """
 
 import torch
-import triton
-import triton.language as tl
 
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import on_gfx1x
+from vllm.triton_utils import tl, triton
 
 torch_dtype: tl.constexpr = torch.float16