prain: fix vLLM on ROCm (no CUDA capability families), smoke test and benchmark scripts

patches/vllm-0.32-rocm-no-cuda-families.patch

@@ -0,0 +1,17 @@
+diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
+index 807e571..bbc642a 100644
+--- a/vllm/platforms/rocm.py
++++ b/vllm/platforms/rocm.py
+@@ -983,6 +983,12 @@ class RocmPlatform(Platform):
+         major, minor = torch.cuda.get_device_capability(device_id)
+         return DeviceCapability(major=major, minor=minor)
+ 
++    @classmethod
++    def is_device_capability_family(cls, capability: int, device_id: int = 0) -> bool:
++        # Families are a CUDA notion, and gfx majors collide with them (gfx1032 reads as
++        # 10.3, i.e. Blackwell sm_103): no CUDA family check may match a ROCm GPU.
++        return False
++
+     @classmethod
+     @with_amdsmi_context
+     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:

scripts/bench.py

@@ -0,0 +1,53 @@
+"""Time to first token and decode speed of an OpenAI-compatible server, as JSON lines.
+
+usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
+Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
+--tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
+the medians.
+"""
+
+import argparse
+import json
+import statistics
+import time
+import uuid
+
+from prain.chat import chat_body, first_model, post, stream_text
+
+PROMPTS = [
+    "Explain how a CPU cache works, step by step, for a beginner.",
+    "Write a short story about a lighthouse keeper who finds a message in a bottle.",
+    "Schreibe eine kurze Anleitung, wie man Brot backt.",
+]
+
+
+def run(base_url, model, prompt, tokens):
+    text = f"[{uuid.uuid4().hex[:8]}] {prompt}"
+    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
+                     temperature=0, ignore_eos=True)
+    start = time.monotonic()
+    stamps = []
+    with post(base_url, "/v1/chat/completions", body, timeout=3600) as response:
+        for _ in stream_text(response):
+            stamps.append(time.monotonic())
+    return {"ttft_s": round(stamps[0] - start, 3),
+            "decode_tok_s": round((tokens - 1) / max(stamps[-1] - stamps[0], 1e-9), 3)}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("base_url")
+    parser.add_argument("--tokens", type=int, default=128)
+    parser.add_argument("--runs", type=int, default=3)
+    args = parser.parse_args()
+    model = first_model(args.base_url)
+    results = []
+    for i in range(args.runs):
+        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens))
+        print(json.dumps(results[-1]), flush=True)
+    print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
+                     | {"median_of": len(results), "tokens": args.tokens}))
+
+
+if __name__ == "__main__":
+    main()

scripts/build.sh

@@ -4,14 +4,16 @@
 #   sdk       venv + AMD's ROCm SDK wheels for $PRAIN_GPU_ARCH
 #   fetch     clone PyTorch and vLLM into build/
 #   torch     PyTorch from source -> build/wheels, installed into the venv
+#   torchvision  torchvision from source -> build/wheels, installed into the venv
 #   vllm      vLLM from source (adds $PRAIN_GPU_ARCH to its HIP arch list), installed editable
 #   prain     this package, installed editable, and its test suite
-# usage: scripts/build.sh STEP...      Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VLLM_REF
+# usage: scripts/build.sh STEP...   Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VISION_REF, PRAIN_VLLM_REF
 set -euo pipefail
 cd "$(dirname "$0")/.."
 ROCM_INDEX=${PRAIN_ROCM_INDEX:-https://stable.repo.amd.com/rocm/whl-next/}
 ROCM_VERSION=${PRAIN_ROCM_VERSION:-10.1.0}
 TORCH_REF=${PRAIN_TORCH_REF:-v2.13.0}
+VISION_REF=${PRAIN_VISION_REF:-v0.28.0}
 VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
 PYTHON=${PRAIN_PYTHON:-python3.13}
 
@@ -90,9 +92,23 @@ torch() {
     env -u LD_LIBRARY_PATH python -c 'import torch; x = torch.ones(4, device="cuda"); print(torch.__version__, torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName, (x * 2).sum().item(), torch.backends.openmp.is_available(), torch.get_num_threads())'
 }
 
+# transformers' Gemma 4 processor imports torchvision, even for text-only serving.
+torchvision() {
+    [ -d build/vision/.git ] || git clone --depth 1 --branch "$VISION_REF" \
+        https://github.com/pytorch/vision build/vision
+    (cd build/vision && env FORCE_CUDA=1 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" MAX_JOBS="$PRAIN_JOBS" \
+        nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v -w "$PRAIN_ROOT/build/wheels" .)
+    python -m pip install --no-deps --force-reinstall build/wheels/torchvision-*.whl
+    env -u LD_LIBRARY_PATH python -c 'import torch, torchvision; print(torchvision.__version__, torch.ops.torchvision.nms(torch.rand(4, 4, device="cuda"), torch.rand(4, device="cuda"), 0.5).device)'
+}
+
 vllm() {
-    local src=build/vllm site
+    local src=build/vllm site p
     site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
+    # prain's fixes to vLLM (patches/vllm-*.patch), each applied once.
+    for p in patches/vllm-*.patch; do
+        git -C $src apply --reverse --check "$PRAIN_ROOT/$p" 2>/dev/null || git -C $src apply "$PRAIN_ROOT/$p"
+    done
     # vLLM lists gfx1030 but not its siblings gfx1031/gfx1032: add ours.
     grep -q "$PRAIN_GPU_ARCH" $src/CMakeLists.txt ||
         sed -i "s/;gfx1030;/;gfx1030;$PRAIN_GPU_ARCH;/" $src/CMakeLists.txt

scripts/smoke.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+# Go/no-go: does this vLLM build run Gemma 4 on this GPU? Serves a 2-layer copy of MODEL_DIR with
+# random weights (port 8000, log in build/smoke.log), then sends a greedy request, the warm-up
+# shape (continue_final_message, priority) and the draft fields (vllm_xargs, cache_salt).
+# usage: scripts/smoke.sh MODEL_DIR [vllm serve args...]
+# Expected: READY, a JSON list of 3 (gibberish) strings, two {"id":"chatcmpl-... lines, STOPPED.
+# RDNA2 needs --attention-backend TRITON_ATTN: ROCM_ATTN's prefill kernel wants 160 KiB of LDS.
+cd "$(dirname "$0")/.." && . scripts/env.sh
+model=$1
+shift
+setsid vllm serve "$model" --served-model-name gemma4 --language-model-only --max-model-len 4096 \
+    --max-num-batched-tokens 1024 --max-num-seqs 2 --kv-cache-memory-bytes 805306368 --port 8000 \
+    --enforce-eager --moe-backend triton --scheduling-policy priority --load-format dummy \
+    --hf-overrides '{"text_config": {"num_hidden_layers": 2, "layer_types": ["sliding_attention", "full_attention"]}}' \
+    "$@" > build/smoke.log 2>&1 < /dev/null &
+server=$!
+for _ in $(seq 120); do
+    curl -sf localhost:8000/health > /dev/null && break
+    kill -0 $server 2> /dev/null || break
+    sleep 5
+done
+curl -sf localhost:8000/health > /dev/null && echo READY || { echo DIED; tail -40 build/smoke.log; exit 1; }
+python scripts/greedy.py http://127.0.0.1:8000 gemma4 | head -c 200; echo
+post() { curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d "$1" | head -c 60; echo; }
+post '{"model": "gemma4", "messages": [{"role": "user", "content": "Why is the sky"}], "max_tokens": 1, "add_generation_prompt": false, "continue_final_message": true, "priority": 1}'
+post '{"model": "gemma4", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 2, "vllm_xargs": {"prain": "draft"}, "cache_salt": "prain-draft"}'
+kill -INT -$server
+while curl -sf localhost:8000/health > /dev/null; do sleep 2; done
+echo STOPPED