[0z5a]

Summary

Add AMD Instinct (ROCm) as a supported execution platform for vllm-rlt without changing the Stage/Loop runtime or model arithmetic. On one MI350X it runs Ouro-1.4B, Ouro-2.6B, Nanbeige4.2-3B and Huginn-0125, and exercises the v0.2.0 feature set: Triton attention, packed prefill via prefill_uva, WCPB, prefix cache, preemption/recovery, asynchronous scheduling, token pipeline, synchronous self-speculation, CUDA Graph replay, static buffers, multi-stream, probed UVA and serving.

The CUDA path is unchanged: ROCm branches are capability-gated, auto selection still prefers FA4/FA3/FA2 on NVIDIA, and the native CUDA build is untouched.

Required human contributor answers

1. What outcome are you trying to achieve?

Make the current v0.2.0 runtime features usable on AMD Instinct with the same public interfaces, validated on the supported model families, in one reviewable PR.

2. What is the highest-leverage variable for that outcome?

The platform-specific decisions that were hard-coded to NVIDIA: attention-backend selection, host-metadata accessibility, and Triton's dot precision. Fixing those unlocks the existing scheduler, graphs, packed-prefill, async and speculation paths without new kernels.

3. How will this work move that variable rather than create busywork?

ROCm auto selection tries Triton then Torch and never imports FlashAttention. UVA is enabled only after a hipHostGetDevicePointer identity probe and has an environment fallback. Triton prefill uses ieee dot precision where ROCm rejects tf32. NVIDIA-only and allocator-specific tests skip with explicit reasons.

4. Who is the customer or user, and what do they actually need?

Users running Ouro, Nanbeige and Huginn on MI300/MI350 who need the same scheduling, cache and graph features without a CUDA stack, and maintainers who need a testable platform boundary.

PR necessity and bugfix evidence

A PR is needed. Before this change ROCm ran only because Triton happened to be the first importable backend after three NVIDIA imports failed; UVA was disabled by not torch.version.hip; Triton prefill failed to compile because ROCm rejects input_precision="tf32"; and the ROCm GPU suite produced 163 native-extension errors plus 13 failures from NVIDIA-only parametrizations. With this PR the suite is green apart from explicit skips, and the new packed/WCPB parity test passes.

Scope

Feature	ROCm status	Evidence
Triton paged attention	Yes	tests/test_attention.py, A/B
FlashAttention 2/3/4	No (NVIDIA-only)	explicit error
Packed prefill / prefill_uva	Opt-in	FP32 greedy matches per-row; BF16 differs by <=0.0156 and is excluded from all
Wavefront chunked prefill	Yes	tests/test_rocm.py parity, A/B
Prefix cache	Yes	GPU suite, A/B
Priority scheduling	Yes	high-priority latency 7.80s → 2.22s, A/B
Incremental KV	Yes	+3.8% at batch 64, A/B
Sparse decode (#157)	Tests pass on stacked branch	137 passed, 3 skipped; tracked in #157
Preemption / recovery	Yes	GPU suite
Async scheduling / execution	Yes	GPU suite, A/B
Async self-speculation	Opt-in; 1.48x at 1.4B (256 out)	A/B, suite
Token pipeline	Opt-in; 2.14x at 1.4B (256 out), 1.63x on 2.6B	tests/test_speculative.py, A/B
Sync self-speculation	Opt-in; 1.41x at 1.4B (256 out), 1.18x on 2.6B	GPU suite, A/B, tuning
CUDA Graph replay	Yes	GPU suite, A/B
UVA host metadata	Yes (probed)	tests/test_tensor_transfer.py, tests/test_rocm.py
PD / KV transfer	Not validated (needs 2 GPUs)	skipped
HTTP serving	Yes	tests/test_serving.py

Model support

Checkpoint	ROCm status
Ouro-1.4B	Validated: suite, A/B, text
Ouro-2.6B	Validated: real-checkpoint generation, shared loader
Ouro-1.4B-Thinking	Validated: real-checkpoint generation (bounded)
Ouro-2.6B-Thinking	Validated: real-checkpoint generation (bounded)
Nanbeige4.2-3B	Validated: suite, throughput, text
Huginn-0125	Validated: graph replay, 32-token generation

Changes

• vllm_rlt/platform.py: is_rocm() / is_cuda() / supports_flash_attention().
• vllm_rlt/attention/backend.py: ROCm auto order is triton, torch; FlashAttention reason is reported.
• vllm_rlt/kernels/triton_attention.py: ieee dot precision on ROCm.
• vllm_rlt/worker/tensors.py: probed UVA with VLLM_RLT_UVA=auto|0|1.
• vllm_rlt/engine/stage_engine.py: platform-neutral graph message.
• Tests: tests/test_rocm.py (platform, UVA, packed/WCPB parity); NVIDIA-only, allocator and multi-GPU skips; resolved graph-row comparison.
• docs/rocm.md, README/docs index/launching links.

Test Plan

vllm-rlt commit: 47c25c3 (branch rocm-support, based on 444638b)

Environment: 1× AMD Instinct MI350X VF (gfx950, 288 GiB HBM), ROCm 7.2, PyTorch 2.14.1+rocm7.2, Triton 3.5.1, Python 3.12, CUDA_VISIBLE_DEVICES=0.

• ruff check . and ruff format --check .: pass.
• CPU suite: pytest -q -m "not gpu" → 909 passed, 46 skipped.
• GPU suite: pytest -q --run-gpu → 1240 passed, 314 skipped.
• Real checkpoints: Ouro-1.4B, Ouro-2.6B, Nanbeige4.2-3B, Huginn-0125.

Test Result

All features (A/B first)

Baseline: same commit, model, device and wheel set, attention_backend="triton", ExecutionConfig() defaults, same ouro_delayed exit mode, VLLM_RLT_UVA=0, and no optional features. There was no earlier ROCm baseline, so all numbers are ROCm-relative; CUDA numbers were not measured on this lease. One warm-up generation is discarded, then the median of three runs is reported (32 output tokens).

Workload	Baseline tok/s	All features tok/s	Speedup
short, batch 1, 64 out	18.4	32.8	1.78x
short, batch 32, 64 out	446	916	2.05x
short, batch 128, 64 out	1148	2909	2.53x
short, batch 256, 64 out	1606	3270	2.04x
short, batch 512, 64 out	1864	3919	2.10x
long, batch 8, 32 out	84.8	171.0	2.02x
long, batch 32, 32 out	153.8	544.4	3.54x
long, batch 64, 32 out	179.8	866.4	4.82x

The 32-token runs are prefill-heavy (80 prompt tokens per request); the 256-token run is decode-heavy. Baseline is 6.3 ms/step versus 3.4 ms/step for all features at batch 128, so the decode path still has host-sync and kernel-count headroom.

all = async scheduling + CUDA Graphs + static buffers + power-of-two padding + UVA + WCPB + prefix cache. Packed prefill is excluded because its BF16 output is not token-identical.

Full-feature sweep

Short context (80-token prompt, 64 output), Ouro-1.4B:

Workload	Baseline	runtime	spec	Best
short b128	1194	2208	2069	1.85x
short b256	1670	4329	2126	2.59x
short b512	2008	4184	2031	2.08x

Long context and checkpoints (1k/4k/8k). runtime = async + graphs + static/pad2 + UVA + WCPB + prefix; spec adds pipeline K4/draft1. Huginn's optimized column is graph replay; Nanbeige has no spec path:

Model	Context	Batch	Out	Baseline	runtime	spec	Best
Ouro-1.4B	1k	32	32	155.7	407.6	761.0	4.89x
Ouro-1.4B	1k	64	32	176.5	661.5	894.2	5.07x
Ouro-1.4B	1k	128	32	177.9	1019.1	998.2	5.73x
Ouro-1.4B	4k	4	128	37.9	84.8	263.6	6.96x
Ouro-1.4B	4k	2	256	28.7	54.7	168.4	5.87x
Ouro-1.4B	8k	2	32	5.0	46.2	107.2	21.4x
Ouro-1.4B	8k	1	128	10.1	27.7	80.3	7.95x
Nanbeige4.2-3B	1k	32	32	167.4	854.7	—	5.11x
Nanbeige4.2-3B	1k	64	32	176.4	1083.7	—	6.14x
Nanbeige4.2-3B	1k	128	32	177.2	1464.8	—	8.27x
Nanbeige4.2-3B	4k	4	64	27.7	199.9	—	7.2x
Nanbeige4.2-3B	8k	2	32	4.2	87.1	—	20.7x
Huginn-0125	1k	1	32	8.8	9.5	—	1.08x
Huginn-0125	1k	2	32	13.6	14.8	—	1.09x
Huginn-0125	1k	4	32	14.0	14.7	—	1.05x
Huginn-0125	1k	8	32	16.5	—	—	—
Huginn-0125	4k	1	32	3.1	3.2	—	1.03x
Huginn-0125	4k	2	32	3.1	—	—	—
Huginn-0125	4k	4	32	3.1	—	—	—
Huginn-0125	4k	8	32	3.1	—	—	—

Huginn is compute-bound by 32 recurrent steps (flat across batch), graph replay adds 3-9%, and 8k is rejected by its context limit. 8k Ouro prefill: 1247 -> 11544 -> 26800 tok/s.

ROCm micro-tuning

Setting	tok/s
default	2798
HIP_FORCE_DEV_KERNARG=1	2781
HSA_ENABLE_INTERRUPT=0	2843
both	2827

PyTorch already reports preferred_blas_library=hipblaslt; forcing it explicitly did not improve throughput, so no rocBLAS fallback is needed. rocprofv3 is available for kernel-vs-wall validation.

Per-feature cold A/B

Single cold run per cell, no warm-up, so it includes JIT and graph-capture cost; directional only. Short is batch 128 output tok/s, long is batch 8 prefill tok/s.

Config	tok/s	vs reference
baseline (non-spec)	970	1.00x
async scheduling	1025	1.06x
CUDA Graphs	986	1.02x
static buffers + pad2	1056	1.09x
UVA	1019	1.05x
WCPB	1402	1.45x
prefix cache	1105	1.14x
all features	2798	2.88x
baseline (spec family, 256 out)	1527	1.00x
sync self-speculation (256 out)	2148	1.41x
token pipeline (256 out)	3270	2.14x
async self-speculation (256 out)	2314	1.48x

Spec rows use 256 output tokens; at 32 output they lose (0.51x/0.95x/0.89x), see the workload-dependent paragraph.

Batch 128, 32 output tokens, one warm-up discarded, median of three; non-spec rows use the same ouro_delayed baseline, spec rows use the legal ouro baseline.

Priority scheduling is positive: with 64 low- and 8 high-priority requests at 16 slots, the high-priority median completion drops from 7.80s (fcfs) to 2.22s (priority); total time moves 7.81s to 8.15s. Incremental KV gives 817 -> 848 tok/s at batch 64 (+3.8%).

Sparse decode (#157) merged with this branch passes 137 passed, 3 skipped on ROCm (attention/batched/fused/sync suites); evidence in #157 (comment). It is tracked in #157 and not part of this PR.

Speculation is workload-dependent. At 32 output tokens the 80-token prefill dominates and it loses; at 256 output tokens the decode-heavy path wins on Ouro-1.4B (baseline 1527 tok/s): pipeline K4/draft1 3270 in the sweep and 2842 in a confirmation run (1.82x), block K2/draft1 2148/2251 (1.41-1.44x), async K2/draft1 2314 (1.48x). On Ouro-2.6B (batch 64, 128 output) pipeline K2/draft1 is 814 vs 500 (1.63x). Greedy token match against baseline is 0.66-0.80 and a 20-question GSM8K check scored 0/20 for both arms, so it stays opt-in and is not enabled by all.

Functional checks

• WCPB produces the same greedy tokens as the default path (tests/test_rocm.py). Packed prefill kernel tests pass, but engine parity on a real Ouro prompt fails and is not claimed.
• Ouro-2.6B: " Paris.\n\nThe final answer is".
• Ouro-1.4B-Thinking: " Paris.\n\nThe capital of France is Paris.".
• Ouro-2.6B-Thinking: " Paris.\n\nThe answer is (A) Paris.\n\n\\boxed{A}".
• Huginn-0125 with --cuda-graphs: " Paris. It is the most populous city in France...", 32 exit depths.

Limitations and follow-ups

• FlashAttention is NVIDIA-only; FA4-specific packed-prefill tests skip.
• Packed prefill is numerically located: FP32 greedy matches per-row exactly, while BF16 attention differs by <=0.0156 and can flip greedy tokens; it stays opt-in and out of all.
• Speculative decoding runs but matched the non-speculative greedy path on 30/32 prompts in a 32-request check; it still needs a dedicated numerics qualification.
• PD/NIXL is not validated (single-GPU instance).
• Sparse decode is tracked in Sparse decode attention: per-token reference lifecycle for async self-speculation #157 (ROCm tests pass on a stacked branch), not part of this PR.
• RFC: [RFC]: Support for current v0.2.0 features on different GPU architectures #172