Another attempt

Signed-off-by: Huy Do <huydhn@gmail.com>

Author: huydhn

.github/scripts/gpt-oss/run_accuracy_checks.sh

@@ -2,6 +2,18 @@
 
 set -eux
 
+# https://docs.vllm.ai/projects/recipes/en/latest/OpenAI/GPT-OSS.html
+if [[ "${DEVICE_TYPE}" == *B200* ]]; then
+  export VLLM_USE_TRTLLM_ATTENTION=1
+  export VLLM_USE_TRTLLM_DECODE_ATTENTION=1
+  export VLLM_USE_TRTLLM_CONTEXT_ATTENTION=1
+  export VLLM_USE_FLASHINFER_MXFP4_BF16_MOE=1
+elif [[ "${DEVICE_NAME}" == *rocm* ]]; then
+  export VLLM_ROCM_USE_AITER=1
+  export VLLM_USE_AITER_UNIFIED_ATTENTION=1
+  export VLLM_ROCM_USE_AITER_MHA=0
+fi
+
 tp=0
 if [[ "${MODEL}" == "openai/gpt-oss-120b" ]]; then
   tp=4

.github/scripts/gpt-oss/run_benchmarks.sh

@@ -2,14 +2,30 @@
 
 set -eux
 
+# https://docs.vllm.ai/projects/recipes/en/latest/OpenAI/GPT-OSS.html
+if [[ "${DEVICE_TYPE}" == *B200* ]]; then
+  export VLLM_USE_TRTLLM_ATTENTION=1
+  export VLLM_USE_TRTLLM_DECODE_ATTENTION=1
+  export VLLM_USE_TRTLLM_CONTEXT_ATTENTION=1
+  export VLLM_USE_FLASHINFER_MXFP4_BF16_MOE=1
+elif [[ "${DEVICE_NAME}" == *rocm* ]]; then
+  export VLLM_ROCM_USE_AITER=1
+  export VLLM_USE_AITER_UNIFIED_ATTENTION=1
+  export VLLM_ROCM_USE_AITER_MHA=0
+fi
+
 pushd vllm-benchmarks/vllm
 cp vllm/benchmarks/lib/utils.py /app/vllm-os-mini/vllm/benchmarks/utils.py || true
 
-if [[ $DEVICE_NAME != 'rocm' ]]; then
+if [[ "${DEVICE_NAME}" != "rocm" ]]; then
   pip install -U openai transformers
   pip install --pre vllm==0.10.1+gptoss \
     --extra-index-url https://wheels.vllm.ai/gpt-oss/ \
     --extra-index-url https://download.pytorch.org/whl/nightly/cu128
+
+  export TORCH_CUDA_ARCH_LIST='8.9 9.0'
+  pip install --no-build-isolation "git+https://github.com/facebookresearch/xformers@v0.0.31" \
+    --extra-index-url https://download.pytorch.org/whl/nightly/cu128
 fi
 
 pip freeze

.github/workflows/gpt-oss-benchmark.yml

@@ -161,29 +161,6 @@ jobs:
             aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws
           fi
 
-          # https://docs.vllm.ai/projects/recipes/en/latest/OpenAI/GPT-OSS.html
-          if [[ "${DEVICE_TYPE}" == *B200* ]]; then
-            export VLLM_USE_TRTLLM_ATTENTION=1
-            export VLLM_USE_TRTLLM_DECODE_ATTENTION=1
-            export VLLM_USE_TRTLLM_CONTEXT_ATTENTION=1
-            export VLLM_USE_FLASHINFER_MXFP4_BF16_MOE=1
-          else
-            export VLLM_USE_TRTLLM_ATTENTION=0
-            export VLLM_USE_TRTLLM_DECODE_ATTENTION=0
-            export VLLM_USE_TRTLLM_CONTEXT_ATTENTION=0
-            export VLLM_USE_FLASHINFER_MXFP4_BF16_MOE=0
-          fi
-
-          if [[ "${DEVICE_NAME}" == *rocm* ]]; then
-            export VLLM_ROCM_USE_AITER=1
-            export VLLM_USE_AITER_UNIFIED_ATTENTION=1
-            export VLLM_ROCM_USE_AITER_MHA=0
-          else
-            export VLLM_ROCM_USE_AITER=0
-            export VLLM_USE_AITER_UNIFIED_ATTENTION=0
-            export VLLM_ROCM_USE_AITER_MHA=0
-          fi
-
           container_name=$(docker run \
             ${GPU_FLAG:-} \
             -e MODEL \
@@ -192,13 +169,6 @@ jobs:
             -e HF_TOKEN \
             -e ENGINE_VERSION \
             -e SAVE_TO_PYTORCH_BENCHMARK_FORMAT \
-            -e VLLM_USE_TRTLLM_ATTENTION \
-            -e VLLM_USE_TRTLLM_DECODE_ATTENTION \
-            -e VLLM_USE_TRTLLM_CONTEXT_ATTENTION \
-            -e VLLM_USE_FLASHINFER_MXFP4_BF16_MOE \
-            -e VLLM_ROCM_USE_AITER \
-            -e VLLM_USE_AITER_UNIFIED_ATTENTION \
-            -e VLLM_ROCM_USE_AITER_MHA \
             --ipc=host \
             --tty \
             --detach \