diff --git a/README.md b/README.md index 00ed4e17..af68f03c 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,9 @@ in mixed prefill/decode batches. ## Getting Started Requires **Python 3.10+** and **PyTorch 2.5+**. For GPU inference, use Linux and -an NVIDIA GPU with a compatible CUDA-enabled PyTorch installation. +an NVIDIA GPU with a compatible CUDA-enabled PyTorch installation. AMD Instinct +GPUs are supported through ROCm; see the [ROCm guide](docs/rocm.md) for the +wheel index, feature matrix and validation results. ```bash # Activate your Python environment first. diff --git a/docs/README.md b/docs/README.md index b0ab7777..0c49cb7d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -27,6 +27,7 @@ an environment to receiving the first generated response. The | [Cache and scheduling features](https://github.com/InferMatrixLab/vllm-rlt/pull/31) | Prefix reuse, incremental KV, priorities, and preemption | | [Prefill/decode disaggregation](https://github.com/InferMatrixLab/vllm-rlt/pull/31) | Single-host GPU worker pools and NIXL transfer | | [HTTP serving](serving.md) | Completions API, streaming, and service lifecycle | +| [AMD ROCm](rocm.md) | Installation, feature matrix, validation results and troubleshooting | | [Profiling](profiling.md) | Core profiler options, Python/HTTP controls, PD collection, and per-rank archives | | [Ouro-Thinking recipe](recipes/ouro_thinking.md) | Shared 1.4B/2.6B setup, JSON/SSE client, budgets and validation limits | | [Accuracy evaluation](accuracy.md) | GSM8K regression setup and comparison methodology | diff --git a/docs/launching.md b/docs/launching.md index 19ebb747..fcee236f 100644 --- a/docs/launching.md +++ b/docs/launching.md @@ -3,7 +3,8 @@ Follow steps 1–8 to serve Ouro-1.4B on one NVIDIA GPU and receive a generated response. Commands work in Bash and fish on Linux; use the activation command for your shell. Run them in the same terminal unless a step explicitly asks -you to open another one. +you to open another one. For AMD Instinct GPUs, use the +[ROCm guide](rocm.md) instead of the `nvidia-smi`/CUDA wheel steps. This guide targets **v0.2.0**. Pin the tag below rather than a moving `main`. For migration from v0.1.0, read the [changelog](../CHANGELOG.md); Python scripts diff --git a/docs/rocm.md b/docs/rocm.md new file mode 100644 index 00000000..7784c29d --- /dev/null +++ b/docs/rocm.md @@ -0,0 +1,259 @@ +# AMD ROCm + +vllm-rlt runs on AMD Instinct GPUs through PyTorch's ROCm build and the Triton +attention backend. The runtime, scheduler, serving entrypoints, asynchronous +scheduling, static buffers, CUDA Graph replay and recurrent-loop features are +shared with the CUDA path. NVIDIA-only packages are not required. + +## Requirements + +- Linux with a ROCm 7.x driver and a `gfx9` Instinct GPU (MI300/MI350 class). +- Python 3.10+ and a ROCm-enabled PyTorch build. The validated stack is + PyTorch 2.14.1 + ROCm 7.2 with Triton 3.5.1 on an MI350X (`gfx950`). +- `rocminfo`/`rocm-smi` must report the GPU and `torch.cuda.is_available()` + must be `True` inside the environment. + +Install PyTorch from the ROCm wheel index, then the project from the +repository root. FlashAttention and NIXL are not needed. + +```bash +python3 -m venv .venv +. .venv/bin/activate +pip install torch --index-url https://download.pytorch.org/whl/rocm7.2 +pip install -e . +python -c "import torch; print(torch.__version__, torch.version.hip, torch.cuda.get_device_name(0))" +``` + +`torch.version.hip` must be set. `vllm-rlt` does not build a CUDA extension, so +there is no `nvcc` step and no separate ROCm package. + +## Run + +Use the Triton backend explicitly, or `auto` (ROCm auto selects Triton and then +Torch; it never tries FlashAttention): + +```bash +python -m vllm_rlt.entrypoints.cli \ + --model /models/Ouro-1.4B \ + --device cuda --dtype bfloat16 \ + --attention-backend triton \ + --prompt "The capital of France is" --max-tokens 32 +``` + +The HTTP server and the Python `LLM` facade take the same arguments +(`attention_backend="triton"` in Python). `CUDA_VISIBLE_DEVICES` is still the +device selector because PyTorch exposes ROCm devices through `torch.cuda`. + +## Feature support + +| Feature | ROCm | Notes | +| --- | --- | --- | +| Ouro-1.4B | Yes | Full validation and A/B benchmark | +| Ouro-2.6B | Yes | Real-checkpoint generation; shared Ouro loader | +| Ouro-1.4B/2.6B-Thinking | Yes | Real-checkpoint generation (bounded) | +| Nanbeige4.2-3B | Yes | Two-loop inference | +| Huginn-0125 | Yes | Synchronous fixed-depth, LAST_EXITED KV, graph replay 32 tokens | +| Triton paged attention | Yes | Decode and prefill | +| Torch reference attention | Yes | Correctness fallback | +| FlashAttention 2/3/4 | No | NVIDIA-only kernels | +| Async scheduling | Yes | `ouro_delayed`, `random_lookahead`, `trace` | +| CUDA Graph replay | Yes | Triton backend, static buffers | +| Multi-stream execution | Yes | `multi_stream` | +| UVA host metadata | Yes (probed) | `hipHostGetDevicePointer` must equal the host address | +| Packed prefill / `prefill_uva` | Opt-in | FP32 greedy matches per-row exactly; BF16 attention differs by <=0.0156 and greedy can diverge, so it stays out of `all` until an accuracy check | +| Wavefront chunked prefill | Yes | `--wavefront-chunked-prefill` | +| Prefix cache | Yes | `enable_prefix_caching=True` | +| Priority scheduling | Yes | `priority=0` is highest; high-priority latency 7.80s → 2.22s | +| Incremental KV | Yes | +3.8% at batch 64 | +| Sparse decode (#157) | ROCm tests pass | `--sparse-decode-attention --sparse-decode-key-retention` on the #157 branch; 137 passed | +| Preemption and recovery | Yes | `enable_preemption=True`, async recovery | +| Token pipeline | Opt-in | **2.14x at 1.4B (256 out)**, 1.63x on Ouro-2.6B; short-output loses | +| Synchronous self-speculation | Opt-in | 1.41x at 1.4B (256 out); 1.18x on Ouro-2.6B | +| Asynchronous self-speculation | Opt-in | 1.48x at 1.4B (256 out); async suite | +| PD (prefill/decode split) | Not validated | Depends on the NIXL/UCX transport | +| HTTP serving | Yes | `/v1/completions` | + +Mapped host metadata (UVA) is probed once per device: Triton receives host +pointers directly, so `hipHostGetDevicePointer` must return the same address. +Set `VLLM_RLT_UVA=0` to force the H2D copy path, or `VLLM_RLT_UVA=1` to skip +the probe. The decision is logged at worker startup. + +## Validation + +The following was measured on one MI350X (288 GiB HBM, `gfx950`) with ROCm 7.2, +PyTorch 2.14.1 and Triton 3.5.1, BF16, greedy decoding, 128 concurrent +requests with 64 output tokens each and an 80-token prompt: + +| Model | Output tokens/s | Requests/s | +| --- | --- | --- | +| Ouro-1.4B | 855 | 13.4 | +| Nanbeige4.2-3B | 1012 | 15.8 | + +Configuration: `--attention-backend triton --async-scheduling --static-buffers +--pad-to-power-of-two --cuda-graphs --cuda-graph-max-graphs 64 --max-num-seqs +128 --max-num-batched-tokens 16384`. Numbers are one run, not a benchmark +contract; the first request includes Triton compilation. + +### Baseline and all-features result + +The baseline is the same commit, model, device and wheel set with +`attention_backend="triton"`, `ExecutionConfig()` defaults, no optional +execution features, the same exit mode and `VLLM_RLT_UVA=0`. There was no +earlier ROCm baseline, so every comparison here is ROCm-relative; CUDA numbers +were not measured on this lease. One warm-up `generate` is discarded, then the +median of three runs is reported (same `ouro_delayed` exit mode for both arms): + +| 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** | + +Long-context prefill throughput rises from 2652 to 5350 tok/s at batch 8, +4811 to 17029 at batch 32, and 5626 to 27102 at batch 64. The long-context +high-concurrency regime is where WCPB, async scheduling and graph replay +compound. + +`all` enables async scheduling, CUDA Graphs, static buffers, power-of-two +padding, UVA, WCPB and prefix cache. Packed prefill is excluded because its +engine parity is still failing (see below). The gain is graph replay plus +batching; short runs are prefill-heavy, so the 256-output decode run is the +conservative 1.72x result (baseline 6.3 ms/step, all 3.4 ms/step at batch 128). + +Micro-tuning on the same `all` configuration (batch 128, no packed): + +| 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 used. + +### 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 (it rejects async/spec); 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 its 32 recurrent steps (throughput flat across batch), graph replay adds 3-9%, and it rejects 8k prompts (context limit). 8k Ouro prefill goes 1247 -> 11544 -> 26800 tok/s. + +### Per-feature warmed A/B + +Batch 128, 32 output tokens, one warm-up discarded, median of three runs. +Non-spec configs use `baseline_delayed` (970 tok/s) as the reference; spec +configs use the `ouro` baseline (994 tok/s), which is the only legal exit mode +for them. + +| 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 | + +The spec rows are the decode-heavy 256-output case. At 32 output tokens the same configs lose (0.51x / 0.95x / 0.89x) because prefill dominates; both directions are kept in the workload-dependent section below. + +Priority and incremental KV are positive: with 64 low- and 8 high-priority +requests at 16 slots, high-priority median completion drops from 7.80s (fcfs) +to 2.22s (priority) while total time moves 7.81s to 8.15s; incremental KV gives +817 -> 848 tok/s at batch 64 (+3.8%). + +Speculation is **workload-dependent**, not simply negative. At 32 output +tokens the 80-token prefill dominates and it loses (best 0.95x); with 256 +output tokens at batch 128 the decode-heavy path wins on Ouro-1.4B: + +| Config | tok/s | vs baseline | +| --- | ---: | ---: | +| baseline | 1527 | 1.00x | +| pipeline K4/draft1 | 3270 | **2.14x** | +| pipeline K6/draft1 | 3249 | 2.13x | +| pipeline K2/draft1 | 2365 | 1.55x | +| pipeline K4/draft2 | 2254 | 1.48x | +| pipeline K2/draft2 | 2251 | 1.47x | +| block K2/draft1 | 2148 | **1.41x** | +| pipeline K1/draft1 | 1905 | 1.25x | +| async K2/draft1 | 2314 | 1.48x | + +A confirmation run measured pipeline K4/draft1 at 2842 tok/s (1.82x), block +K2/draft1 at 2251 (1.44x), and async K2/draft1 at 2314 (1.48x), with greedy +token match 0.66-0.74 against baseline. + +On Ouro-2.6B (batch 64, 128 output): pipeline K2/draft1 814 vs 500 (1.63x), +K4/draft1 782 (1.56x), block K2/draft1 591 (1.18x). Greedy token match against +the non-speculative baseline is 0.65-0.80, and a 20-question GSM8K check scored +0/20 for both arms, so there is no accuracy signal yet; speculation stays +opt-in and is not a default. + +### Not done on ROCm + +- **Packed prefill (`prefill_uva`)**: root cause located as BF16 accumulation, + not metadata: FP32 greedy output matches per-row exactly, while BF16 + per-call attention differs by <=0.0156 and can flip greedy tokens. It stays + opt-in until an accuracy evaluation, and is excluded from `all`. +- **PD/KV transfer**: needs two GPUs and a ROCm NIXL/UCX transport. +- **KV quantization / INT4 residual cache**: not present on the current main + runtime; tracked by RFC #156. Sparse decode from #157 passes its ROCm suites + on a stacked branch but is not shipped here. +- **Profiling flags**: `--profile` / `--profile-activities` are untested on + ROCm; `rocprofv3` is available for external traces. +- **FlashAttention**: NVIDIA-only by construction. +- **Multi-node**: out of scope for this device. + +The CPU regression suite passes unchanged on ROCm. GPU tests use the same +`--run-gpu` flag; FlashAttention, FA4-specific packed-prefill and multi-GPU PD tests +skip because those implementations or devices are unavailable. + +## Troubleshooting + +- `torch.version.hip` is `None`: a CUDA PyTorch wheel is installed. Reinstall + with the ROCm index above. +- `HIP error: invalid device function` or an illegal memory access inside + Triton: verify the GPU architecture is supported by the installed ROCm and + Triton versions, and that `torch.cuda.get_device_name(0)` reports the device. +- `flash_attn currently requires an NVIDIA CUDA device`: use + `--attention-backend triton`; FlashAttention cannot run on ROCm. diff --git a/tests/conftest.py b/tests/conftest.py index f65e8a42..8250d0a4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,17 @@ def pytest_collection_modifyitems(config, items): for item in items: if "gpu" in item.keywords and not enabled: item.add_marker(pytest.mark.skip(reason="set CUDA_VISIBLE_DEVICES and use --run-gpu")) + continue + callspec = getattr(item, "callspec", None) + if ( + enabled + and torch.version.hip + and callspec is not None + and "flash_attn" in str(callspec.params) + ): + # FlashAttention packages have no ROCm kernels; the Triton backend + # covers these parametrizations on AMD hardware. + item.add_marker(pytest.mark.skip(reason="FlashAttention is NVIDIA-only")) @pytest.fixture(autouse=True, scope="session") diff --git a/tests/test_activation.py b/tests/test_activation.py index 0b877a0b..bc0234af 100644 --- a/tests/test_activation.py +++ b/tests/test_activation.py @@ -135,6 +135,8 @@ def import_module(name, *args, **kwargs): @pytest.fixture(autouse=True) def require_native_extension_for_gpu_tests(request): if request.node.get_closest_marker("gpu"): + if torch.version.hip is not None: + pytest.skip("native CUDA SiLU-and-Mul is not built on ROCm") from vllm_rlt.kernels._extension import load_extension load_extension() diff --git a/tests/test_cdb_runtime.py b/tests/test_cdb_runtime.py index 8b345dcb..912caa3b 100644 --- a/tests/test_cdb_runtime.py +++ b/tests/test_cdb_runtime.py @@ -8,6 +8,7 @@ from tests.helpers import LocalLLM as LLM from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.config import resolve_cuda_graph_rows from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.core.memory import budget_blocks from vllm_rlt.engine.llm_engine import LLMEngine @@ -597,7 +598,9 @@ def test_cuda_dynamic_async_respects_buffers_and_matches_sync(layout, multi_stre expected = LLM(base, **common).generate(prompts, params) execution = ExecutionConfig(async_scheduling=True, multi_stream=multi_stream) actual = LLM(base, execution_config=execution, **common) - assert actual.engine.execution_config == execution + assert actual.engine.execution_config == resolve_cuda_graph_rows( + execution, common["scheduler_config"] + ) assert actual.engine.model_runner.requestspaces == {} assert actual.engine.model_runner.states is actual.engine.model_runner.async_state.hidden result = actual.generate(prompts, params) @@ -634,6 +637,7 @@ def test_cuda_graph_reserve_reduces_budget_with_fixed_memory_inputs(monkeypatch, @pytest.mark.gpu +@pytest.mark.skipif(torch.version.hip is not None, reason="CUDA allocator residue accounting") def test_cuda_memory_plan_reclaims_allocator_residue_on_recreation(): import gc diff --git a/tests/test_flash_attention.py b/tests/test_flash_attention.py index cb2c70bc..652a57c8 100644 --- a/tests/test_flash_attention.py +++ b/tests/test_flash_attention.py @@ -10,6 +10,10 @@ from vllm_rlt.kernels.paged_attention import torch_paged_attention from vllm_rlt.models import OuroForCausalLM +pytestmark = pytest.mark.skipif( + torch.version.hip is not None, reason="FlashAttention is NVIDIA-only" +) + @pytest.mark.parametrize( "sm,expected", [((8, 0), 2), ((8, 9), 2), ((9, 0), 3), ((10, 0), 4), ((10, 3), 4)] diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py index 6f580135..202935a2 100644 --- a/tests/test_input_preparation.py +++ b/tests/test_input_preparation.py @@ -280,6 +280,7 @@ def reject(*args, **kwargs): @pytest.mark.gpu +@pytest.mark.skipif(torch.version.hip is not None, reason="packed prefill is FA4-specific") @pytest.mark.parametrize("use_uva", [False, True]) @pytest.mark.parametrize("single_query", [False, True]) def test_packed_metadata_supports_query_lengths_and_per_request_depths(use_uva, single_query): diff --git a/tests/test_ouro.py b/tests/test_ouro.py index e7f29cd3..63adee7a 100644 --- a/tests/test_ouro.py +++ b/tests/test_ouro.py @@ -36,7 +36,18 @@ def test_cuda_mlp_matches_torch_expression(dtype, rows): @pytest.mark.gpu -@pytest.mark.parametrize("native_available", [True, False]) +@pytest.mark.parametrize( + "native_available", + [ + pytest.param( + True, + marks=pytest.mark.skipif( + torch.version.hip is not None, reason="native CUDA extension is NVIDIA-only" + ), + ), + False, + ], +) def test_cuda_mlp_selects_native_or_missing_extension_fallback(native_available): # Dispatcher registrations last for the process lifetime. Exercise an absent # extension in a fresh process rather than changing an already loaded backend. diff --git a/tests/test_pd.py b/tests/test_pd.py index bb1554c5..12a22403 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -608,6 +608,7 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling): @pytest.mark.gpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="PD needs two visible GPUs") @pytest.mark.parametrize("p_cache,d_cache", [(True, True), (True, False), (False, True)]) @pytest.mark.parametrize("graph", [False, True]) def test_pd_prefix_reuse_reduces_transfers_and_preserves_outputs(p_cache, d_cache, graph): diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py index c31d4aeb..67bdd3f0 100644 --- a/tests/test_prefix_growth.py +++ b/tests/test_prefix_growth.py @@ -157,6 +157,7 @@ def test_lossless_preemption_preserves_looped_history(): @pytest.mark.gpu +@pytest.mark.skipif(torch.version.hip is not None, reason="packed prefill is FA4-specific") @pytest.mark.parametrize("graph", [False, True]) @pytest.mark.parametrize("use_uva", [False, True]) def test_fa4_prefill_uva_prefix_growth_and_bank_reuse(graph, use_uva): @@ -236,6 +237,7 @@ def test_priority_preempts_at_safe_boundary(): @pytest.mark.gpu +@pytest.mark.skipif(torch.version.hip is not None, reason="packed prefill is FA4-specific") @pytest.mark.parametrize("use_uva", [False, True]) def test_prefill_uva_metadata_matches_reference_and_waits_for_consumer(use_uva): from vllm_rlt.worker.buffers import PackedMetadataBuffer diff --git a/tests/test_rocm.py b/tests/test_rocm.py new file mode 100644 index 00000000..f288148b --- /dev/null +++ b/tests/test_rocm.py @@ -0,0 +1,86 @@ +"""ROCm platform behavior that does not require AMD hardware.""" + +import pytest +import torch + +from vllm_rlt.attention import create_backend +from vllm_rlt.platform import is_cuda, is_rocm, supports_flash_attention +from vllm_rlt.worker.tensors import use_uva_on + + +@pytest.fixture +def rocm(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "7.2.0") + monkeypatch.setattr(torch.version, "cuda", None) + + +def test_platform_flags(rocm): + assert is_rocm() + assert not is_cuda() + assert not supports_flash_attention() + + +def test_auto_prefers_triton_on_rocm(rocm): + selected = create_backend("auto", torch.device("cuda"), torch.bfloat16, 128, 16) + assert selected.info["selected_backend"] == "triton" + assert "FlashAttention is NVIDIA-only" in selected.info["selection_reason"] + + +def test_explicit_flash_rejected_on_rocm(rocm): + with pytest.raises(ValueError, match="NVIDIA CUDA"): + create_backend("flash_attn", torch.device("cuda"), torch.bfloat16, 128, 16) + + +def test_uva_override_and_cpu_fallback(monkeypatch): + assert not use_uva_on(torch.device("cpu")) + monkeypatch.setenv("VLLM_RLT_UVA", "0") + assert not use_uva_on(torch.device("cuda", 0)) + monkeypatch.setenv("VLLM_RLT_UVA", "1") + assert use_uva_on(torch.device("cuda", 0)) is torch.cuda.is_available() + + +@pytest.mark.gpu +def test_packed_prefill_and_wcpb_match_default(): + from tests.helpers import tiny_ouro_config + from vllm_rlt import ( + LLM, + CacheConfig, + ExecutionConfig, + ExitConfig, + SamplingParams, + SchedulerConfig, + ) + from vllm_rlt.models import OuroForCausalLM + + torch.manual_seed(9) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16) + prompts = [[2] * 16, [3] * 23] + params = [SamplingParams(max_tokens=3, ignore_eos=True, temperature=0.0)] * 2 + + def run(*, execution=None, scheduler=None): + with LLM( + model, + device="cuda", + dtype=torch.bfloat16, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=64, block_size=16), + scheduler_config=scheduler + or SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=64, prefill_chunk_size=8), + execution_config=execution or ExecutionConfig(), + exit_config=ExitConfig("ouro"), + ) as llm: + return [o.token_ids for o in llm.generate(prompts, params)] + + baseline = run() + assert run(execution=ExecutionConfig(prefill_uva=True, use_uva=True)) == baseline + assert ( + run( + scheduler=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=64, + prefill_chunk_size=8, + wavefront_chunked_prefill=True, + ) + ) + == baseline + ) diff --git a/tests/test_tensor_transfer.py b/tests/test_tensor_transfer.py index 61dd0a90..fe9af7dc 100644 --- a/tests/test_tensor_transfer.py +++ b/tests/test_tensor_transfer.py @@ -388,9 +388,11 @@ def test_uva_configuration_and_capability(monkeypatch): assert parser.parse_args([]).use_uva args = parser.parse_args(["--no-uva", "--prefill-uva"]) assert not args.use_uva and args.prefill_uva - monkeypatch.setattr(torch.version, "hip", None) - assert use_uva_on("cuda") - assert not use_uva_on("cuda", False) assert not use_uva_on("cpu") - monkeypatch.setattr(torch.version, "hip", "test") + monkeypatch.setenv("VLLM_RLT_UVA", "0") assert not use_uva_on("cuda") + monkeypatch.setenv("VLLM_RLT_UVA", "1") + assert use_uva_on("cuda") is torch.cuda.is_available() + monkeypatch.delenv("VLLM_RLT_UVA") + assert use_uva_on("cuda") is torch.cuda.is_available() + assert not use_uva_on("cuda", False) diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..fcf6d9d7 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -9,6 +9,7 @@ from vllm_rlt.kernels.flash_attention import FLASH_BACKENDS, FlashPagedAttention from vllm_rlt.kernels.paged_attention import TritonPagedAttention, torch_paged_attention +from vllm_rlt.platform import is_rocm @dataclass(frozen=True) @@ -46,6 +47,8 @@ def generation(self): _TORCH_CAPABILITIES = BackendCapabilities() +# Triton's packed prefill opts in through required_capabilities (prefill_uva); +# it is not the default because FA4 owns packed prefill on CUDA. _TRITON_CAPABILITIES = BackendCapabilities( cuda_graphs=True, async_scheduling=True, device_context_lengths=True ) @@ -100,10 +103,12 @@ def create_backend( ) -> AttentionBackend: """Explicit choices fail without fallback; auto tries FA4/3/2, Triton, then Torch. - Each candidate validates hardware/configuration and imports its implementation. - Only compatibility and import failures permit auto to try the next candidate; - kernel execution errors never trigger a runtime fallback. Execution requirements - filter candidates before memory profiling or device-buffer allocation. + ROCm skips the NVIDIA-only FlashAttention candidates and tries Triton then + Torch. Each candidate validates hardware/configuration and imports its + implementation. Only compatibility and import failures permit auto to try the + next candidate; kernel execution errors never trigger a runtime fallback. + Execution requirements filter candidates before memory profiling or + device-buffer allocation. """ validate_backend_name(backend) candidates = _AUTO_PRIORITY if backend == "auto" else (backend,) @@ -111,6 +116,9 @@ def create_backend( if backend == "auto" and device.type != "cuda": candidates = ("torch",) skipped.append(f"{device.type.upper()} uses Torch reference attention") + elif backend == "auto" and is_rocm(): + candidates = ("triton", "torch") + skipped.append("ROCm uses Triton or Torch attention; FlashAttention is NVIDIA-only") for candidate in candidates: try: implementation = _create_explicit(candidate, device, dtype, head_dim, block_size) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..e1c06983 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -122,7 +122,7 @@ def __init__( cache_config, memory_reserve_bytes=cache_config.memory_reserve_bytes + scratch ) if self.execution_config.cuda_graphs and parameter.device.type != "cuda": - raise ValueError("CUDA graphs require CUDA with Triton or FlashAttention") + raise ValueError("CUDA graphs require CUDA or ROCm with Triton or FlashAttention") if self.execution_config.cuda_graphs and speculative_config is not None: rows = min( scheduler_config.max_num_seqs * (speculative_config.num_speculative_tokens + 1), @@ -139,6 +139,7 @@ def __init__( scheduler_config.max_num_seqs, speculative_config.num_speculative_tokens, ) + if self.execution_config.async_scheduling: if not self.async_speculative and self.exit_config.mode not in ( "ouro_delayed", diff --git a/vllm_rlt/kernels/triton_attention.py b/vllm_rlt/kernels/triton_attention.py index f1953a92..c784487f 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -346,7 +346,8 @@ def paged_prefill_attention( BLOCK_M=block_m, BLOCK_N=block_n, # FP32 callers (reference checks) need exact products, not TF32. - PRECISION="ieee" if q.dtype == torch.float32 else "tf32", + # ROCm Triton rejects "tf32" even for BF16/FP16 inputs; only "ieee" is accepted. + PRECISION="ieee" if torch.version.hip or q.dtype == torch.float32 else "tf32", max_q_blocks=q_blocks, num_warps=4, num_stages=2, diff --git a/vllm_rlt/platform.py b/vllm_rlt/platform.py new file mode 100644 index 00000000..9102e9a0 --- /dev/null +++ b/vllm_rlt/platform.py @@ -0,0 +1,18 @@ +"""Accelerator platform facts used by backend selection and storage policies.""" + +import torch + + +def is_rocm() -> bool: + """Whether PyTorch targets AMD ROCm (HIP) instead of NVIDIA CUDA.""" + return torch.version.hip is not None + + +def is_cuda() -> bool: + """Whether PyTorch targets NVIDIA CUDA.""" + return torch.version.cuda is not None + + +def supports_flash_attention() -> bool: + """Whether the FlashAttention packages can run on this build.""" + return is_cuda() and not is_rocm() diff --git a/vllm_rlt/worker/tensors.py b/vllm_rlt/worker/tensors.py index 503dee5c..ff811184 100644 --- a/vllm_rlt/worker/tensors.py +++ b/vllm_rlt/worker/tensors.py @@ -7,10 +7,13 @@ These objects never synchronize or decide when their storage may be overwritten. Callers must acquire the owning buffer before writes and retire its final consumer -before reset/reuse. CUDA graph inputs and hot hidden/KV states stay device-resident. +before reset/reuse. Graph inputs and hot hidden/KV states stay device-resident. """ +import ctypes +import logging import math +import os from functools import cache, cached_property import torch @@ -20,9 +23,49 @@ _DTYPES = (torch.int32, torch.int64, torch.float16, torch.bfloat16, torch.float32) +def _hip_uva_identity(device_index: int) -> bool: + """Whether ROCm returns the same device address for mapped pinned memory. + + Triton receives the host pointer directly, so a distinct device alias is not + usable without passing both pointers through the runtime. hipHostGetDevicePointer + is a host-side query: it does not launch a kernel or fault on unmapped memory. + """ + try: + with torch.cuda.device(device_index): + probe = torch.zeros(1, dtype=torch.int32, pin_memory=True) + library = ctypes.CDLL("libamdhip64.so") + device_pointer = ctypes.c_void_p() + status = library.hipHostGetDevicePointer( + ctypes.byref(device_pointer), ctypes.c_void_p(probe.data_ptr()), 0 + ) + return status == 0 and device_pointer.value == probe.data_ptr() + except (OSError, RuntimeError): + return False + + +@cache +def _uva_available(device_index: int) -> bool: + available = device_index < torch.cuda.device_count() + if available and torch.version.hip: + available = _hip_uva_identity(device_index) + logging.getLogger(__name__).info( + "Mapped host metadata on device %d: %s", device_index, "enabled" if available else "off" + ) + return available + + def use_uva_on(device, enabled=True): - """Mapped host reads are supported by the NVIDIA CUDA path.""" - return enabled and torch.device(device).type == "cuda" and not torch.version.hip + """Mapped host reads are probed once per device; VLLM_RLT_UVA=0 forces them off.""" + target = torch.device(device) + if not enabled or target.type != "cuda" or not torch.cuda.is_available(): + return False + mode = os.environ.get("VLLM_RLT_UVA", "auto").strip().lower() + if mode in ("0", "false", "off", "no"): + return False + if mode in ("1", "true", "on", "yes"): + return True + index = target.index if target.index is not None else torch.cuda.current_device() + return _uva_available(index) @cache @@ -36,7 +79,7 @@ def copy_tensor(source, target, *, secondary=None, use_uva=True): """Copy contiguous values directly to final storage, optionally casting twice. Source and outputs must stay alive through the caller's completion event. - D2H deliberately uses CUDA's async copy path rather than SM writes to host. + D2H deliberately uses the device async copy path rather than SM writes to host. """ if source.numel() != target.numel() or not source.is_contiguous() or not target.is_contiguous(): raise ValueError("copy_tensor requires equal-sized contiguous tensors") @@ -56,10 +99,8 @@ def copy_tensor(source, target, *, secondary=None, use_uva=True): raise ValueError("cross-device transfers require an explicit connector") if source.numel() == 0: return target - direct = ( - target.is_cuda - and not torch.version.hip - and (source.is_cuda or (use_uva_on(target.device, use_uva) and source.is_pinned())) + direct = target.is_cuda and ( + source.is_cuda or (use_uva_on(target.device, use_uva) and source.is_pinned()) ) if direct: from vllm_rlt.kernels.tensor_transfer import tensor_transfer_kernel