add triton_key

benchmarks/benchmark_prefix_caching.py

@@ -38,7 +38,7 @@ from transformers import PreTrainedTokenizerBase
 from vllm import LLM, SamplingParams
 from vllm.engine.arg_utils import EngineArgs
 from vllm.utils import FlexibleArgumentParser
-import triton
+# import triton
 
 
 try:
@@ -46,23 +46,23 @@ try:
 except ImportError:
     from backend_request_func import get_tokenizer
 
-triton_version = triton.__version__
-if triton_version.startswith("2.1"):
-    from triton.common.backend import compute_core_version_key
-elif triton_version.startswith("3.0"):
-    from triton.compiler.compiler import triton_key
-else:
-    print(f"TRITON version {triton_version} is not specifically handled.")
+# triton_version = triton.__version__
+# if triton_version.startswith("2.1"):
+#     from triton.common.backend import compute_core_version_key
+# elif triton_version.startswith("3.0"):
+#     from triton.compiler.compiler import triton_key
+# else:
+#     print(f"TRITON version {triton_version} is not specifically handled.")
 
 
 PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as fellows. You need to answer my question about the table.\n# Table\n|Opening|Opening|Sl. No.|Film|Cast|Director|Music Director|Notes|\n|----|----|----|----|----|----|----|----|\n|J A N|9|1|Agni Pushpam|Jayabharathi, Kamalahasan|Jeassy|M. K. Arjunan||\n|J A N|16|2|Priyamvada|Mohan Sharma, Lakshmi, KPAC Lalitha|K. S. Sethumadhavan|V. Dakshinamoorthy||\n|J A N|23|3|Yakshagaanam|Madhu, Sheela|Sheela|M. S. Viswanathan||\n|J A N|30|4|Paalkkadal|Sheela, Sharada|T. K. Prasad|A. T. Ummer||\n|F E B|5|5|Amma|Madhu, Srividya|M. Krishnan Nair|M. K. Arjunan||\n|F E B|13|6|Appooppan|Thikkurissi Sukumaran Nair, Kamal Haasan|P. Bhaskaran|M. S. Baburaj||\n|F E B|20|7|Srishti|Chowalloor Krishnankutty, Ravi Alummoodu|K. T. Muhammad|M. S. Baburaj||\n|F E B|20|8|Vanadevatha|Prem Nazir, Madhubala|Yusufali Kechery|G. Devarajan||\n|F E B|27|9|Samasya|Madhu, Kamalahaasan|K. Thankappan|Shyam||\n|F E B|27|10|Yudhabhoomi|K. P. Ummer, Vidhubala|Crossbelt Mani|R. K. Shekhar||\n|M A R|5|11|Seemantha Puthran|Prem Nazir, Jayabharathi|A. B. Raj|M. K. Arjunan||\n|M A R|12|12|Swapnadanam|Rani Chandra, Dr. Mohandas|K. G. George|Bhaskar Chandavarkar||\n|M A R|19|13|Thulavarsham|Prem Nazir, sreedevi, Sudheer|N. Sankaran Nair|V. Dakshinamoorthy||\n|M A R|20|14|Aruthu|Kaviyoor Ponnamma, Kamalahasan|Ravi|G. Devarajan||\n|M A R|26|15|Swimming Pool|Kamal Haasan, M. G. Soman|J. Sasikumar|M. K. Arjunan||\n\n# Question\nWhat' s the content in the (1,1) cells\n"  # noqa: E501
 
 
 def test_prefix(llm=None, sampling_params=None, prompts=None):
-    if triton_version.startswith("2.1"):
-        version_key = compute_core_version_key()
-    if triton_version.startswith("3.0"):
-        version_key = triton_key()
+    # if triton_version.startswith("2.1"):
+    #     version_key = compute_core_version_key()
+    # if triton_version.startswith("3.0"):
+    #     version_key = triton_key()
     start_time = time.time()
 
     llm.generate(prompts, sampling_params=sampling_params)

vllm/attention/backends/rocm_flash_attn.py

@@ -5,6 +5,8 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
 
 import torch
+import triton
+from triton.compiler.compiler import triton_key
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
@@ -781,6 +783,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
             else:
                 # prefix-enabled attention -
                 # not applicable for encoder-only models
+                version_key = triton_key()
                 if self.attn_type != AttentionType.ENCODER_ONLY:
                     output[:
                            num_prefill_tokens] = PagedAttention.forward_prefix(