fix amd Examples Test

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

Author: Concurrensee

.buildkite/test-amd.yaml

@@ -434,7 +434,7 @@ steps:
 
 - label: Examples Test # 30min
   timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
+  mirror_hardwares: [amdexperimental, amdproduction]
   agent_pool: mi325_1
   # grade: Blocking
   working_dir: "/vllm-workspace/examples"

examples/offline_inference/basic/embed.py

@@ -1,9 +1,11 @@
 # 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 +22,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,11 @@
 # 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 +22,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 = [