fix amd Examples Test

Signed-off-by: Yida Wu <yida.wu@amd.com>

Author: Concurrensee

examples/offline_inference/basic/embed.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
+from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -20,6 +23,8 @@ def parse_args():
 
 
 def main(args: Namespace):
+    if current_platform.is_rocm():
+        os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"
     # Sample prompts.
     prompts = [
         "Hello, my name is",

examples/offline_inference/basic/score.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
+from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -20,6 +23,8 @@ def parse_args():
 
 
 def main(args: Namespace):
+    if current_platform.is_rocm():
+        os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"
     # Sample prompts.
     text_1 = "What is the capital of France?"
     texts_2 = [