diff --git a/patches/vllm-0.32-rocm-no-cuda-families.patch b/patches/vllm-0.32-rocm-no-cuda-families.patch new file mode 100644 index 0000000..09220d6 --- /dev/null +++ b/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: diff --git a/scripts/bench.py b/scripts/bench.py new file mode 100644 index 0000000..fc556f6 --- /dev/null +++ b/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() diff --git a/scripts/build.sh b/scripts/build.sh index 90539e2..2c370d5 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -103,8 +103,12 @@ torchvision() { } 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 diff --git a/scripts/smoke.sh b/scripts/smoke.sh new file mode 100755 index 0000000..d2a2247 --- /dev/null +++ b/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