diff --git a/README.md b/README.md index 00ed4e17..4fc95ae5 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,9 @@ in mixed prefill/decode batches. - **Configurable execution:** Torch reference execution, GPU attention backends, CUDA Graphs, and optional asynchronous/speculative paths. An optional native CUDA SiLU-and-Mul operator accelerates the Ouro activation path, with a Torch - fallback; it is not a general model-speedup claim. + fallback; it is not a general model-speedup claim. FlashInfer FA2 paged decode + is an explicit opt-in backend for Ada (SM89) GPUs with CUDA Graphs and async + scheduling; see [the FA2 guide](docs/flashinfer-fa2.md). - **Offline and online use:** Python API, CLI, and a streaming OpenAI-compatible **completions** endpoint. This is not a chat-completions API. diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index d1ecc7f0..ddeba495 100644 --- a/benchmarks/gsm8k.py +++ b/benchmarks/gsm8k.py @@ -559,6 +559,8 @@ def run(args): import torch from transformers import AutoTokenizer + if getattr(args, "cuda_graphs", False) and args.backend != "native": + raise ValueError("--cuda-graphs applies to the native backend only") protocol = json.loads(Path(args.protocol).read_text()) if ( "baseline" in protocol @@ -587,10 +589,14 @@ def run(args): raise ValueError("Checkpoint depth differs from the protocol") if not os.environ.get("CUDA_VISIBLE_DEVICES") or torch.cuda.device_count() != 1: raise RuntimeError("Run with one scheduler-assigned GPU") + attention_backend = getattr(args, "attention_backend", "triton") + cuda_graphs = getattr(args, "cuda_graphs", False) output = Path(args.output) output.mkdir(parents=True, exist_ok=False) metadata = { "backend": args.backend, + "attention_backend": attention_backend, + "cuda_graphs": cuda_graphs, "protocol_sha256": file_digest(args.protocol), "source": source(), "cuda_visible_devices": os.environ["CUDA_VISIBLE_DEVICES"], @@ -609,7 +615,15 @@ def run(args): close_id = thinking_close_id(tokenizer) if scorer == SCORER_THINKING else None task = make_task(protocol["task_config"]) start = time.monotonic() - generator = Generator(args.backend, str(model), tokenizer, protocol["max_length"], exit_policy) + generator = Generator( + args.backend, + str(model), + tokenizer, + protocol["max_length"], + exit_policy, + attention_backend=attention_backend, + cuda_graphs=cuda_graphs, + ) load_seconds = time.monotonic() - start rows = [] try: @@ -785,6 +799,12 @@ def main(): ) p = commands.add_parser("run") p.add_argument("--backend", choices=["transformers", "native"], required=True) + p.add_argument( + "--attention-backend", + choices=["torch", "triton", "flashinfer_fa2"], + default="triton", + help="Native backend's attention implementation", + ) p.add_argument("--protocol", required=True) p.add_argument("--output", required=True) p.add_argument("--exit-mode", choices=["ouro", "ouro_delayed"], default="ouro") @@ -796,6 +816,11 @@ def main(): ) p.add_argument("--min-loops", type=int, help="Required when --exit-threshold is below 1") p.add_argument("--async-scheduling", action="store_true", help="Native ouro_delayed only") + p.add_argument( + "--cuda-graphs", + action="store_true", + help="Native recurrent-core CUDA Graphs; recorded, not part of the baseline fingerprint", + ) p = commands.add_parser("compare") p.add_argument("--transformers", "--reference", dest="transformers", required=True) p.add_argument("--native", "--candidate", dest="native", required=True) diff --git a/benchmarks/gsm8k_backends.py b/benchmarks/gsm8k_backends.py index 15d57165..8cc73fda 100644 --- a/benchmarks/gsm8k_backends.py +++ b/benchmarks/gsm8k_backends.py @@ -101,7 +101,16 @@ def _after_model(self, module, inputs, output): class Generator: - def __init__(self, backend, model_path, tokenizer, max_length, exit_policy): + def __init__( + self, + backend, + model_path, + tokenizer, + max_length, + exit_policy, + attention_backend="triton", + cuda_graphs=False, + ): self.backend, self.tokenizer, self.exit = backend, tokenizer, exit_policy # Full depth, taken from the protocol (the checkpoint's total_ut_steps). loops = exit_policy["max_loops"] @@ -138,12 +147,14 @@ def __init__(self, backend, model_path, tokenizer, max_length, exit_policy): check_depth(self.model.config, loops) self.engine = StageEngine( self.model, - attention_backend="triton", + attention_backend=attention_backend, # last_exited KV keeps one plane per loop. cache_config=CacheConfig(num_blocks=loops * ((max_length + 15) // 16)), scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=128), exit_config=ExitConfig(exit_policy["mode"]), - execution_config=ExecutionConfig(async_scheduling=exit_policy["async_scheduling"]), + execution_config=ExecutionConfig( + async_scheduling=exit_policy["async_scheduling"], cuda_graphs=cuda_graphs + ), ) # Observe the existing sampling boundary without changing arithmetic. self._finite_hook = self.model.lm_head.register_forward_hook( diff --git a/docs/README.md b/docs/README.md index b0ab7777..d830cb29 100644 --- a/docs/README.md +++ b/docs/README.md @@ -24,6 +24,7 @@ an environment to receiving the first generated response. The | [Self-speculation](async-speculative.md) | Cross-round asynchronous block execution, token pipelines, and acceptance/rejection timelines | | [Asynchronous scheduling](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution | | [FlashAttention](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints | +| [FlashInfer FA2 (Ada)](flashinfer-fa2.md) | SM89 tensor-core paged decode, CUDA Graphs, explicit opt-in, validation scope | | [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 | diff --git a/docs/accuracy.md b/docs/accuracy.md index fd1f9463..65cbb8ad 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -198,6 +198,50 @@ The [Ouro evaluation settings](https://arxiv.org/html/2510.25741v5#A3.T16) do no pin the exact harness revision, demonstrations or token limits, so the settings above are explicit project choices rather than an exact paper reproduction. +## Attention-backend runs + +`benchmarks.gsm8k run` accepts `--attention-backend {torch,triton,flashinfer_fa2}` +for the native backend; the default stays `triton`. `--cuda-graphs` selects the +native recurrent-core CUDA Graphs for that run. Both flags are recorded in the +summary and do not change the frozen baseline fingerprint. + +On one L20 (SM89, torch 2.13.0, CUDA 13.0, FlashInfer 0.7.2.post1, Triton +3.7.1), the default GSM8K-87 protocol (HF baseline 59/87) produced: + +| run | correct | accuracy | gate | +|---|---:|---:|---| +| triton | 61/87 | 70.11% | pass (+2.30 pp) | +| flashinfer_fa2 eager | 59/87 | 67.82% | pass (+0.00 pp) | +| flashinfer_fa2 + CUDA Graphs | 59/87 | 67.82% | pass (+0.00 pp) | + +Eager and graph FA2 are **bit-identical (87/87 token sequences)**. This backend +disables FlashInfer's split-KV: with splitting, a row's reduction order depends +on the planned width and batch composition, which made the earlier eager and +graph runs differ on 18 questions. Without splits, 67/87 sequences match +Triton; four answers flip (203, 426, 819 correct → incorrect; 871 incorrect → +correct). There is no length-limited runaway (mean 106.7 tokens; eager 717.8 s, +graphs 279.6 s). The no-split arithmetic reproduces the earlier split-KV graph +run on all 87 questions; the earlier split-KV eager run scored 60/87. + +The `flashinfer_fa2` backend runs FlashInfer's tensor-core, CUDA-Graph- +compatible decode wrapper; the FA2 guide records why the CUDA-core path is not +graph-safe on SM89. CUDA Graphs remove most per-token launch overhead on this +small model for both attention backends: a matched 64-token fixed-work run at +concurrency 1/8/32 went from 5.49/5.55/6.32 s (FA2 eager) to 1.59/1.84/2.24 s +(FA2 graphs), and from 5.20/5.34/6.26 s (Triton eager) to 1.52/1.76/2.20 s +(Triton graphs). Triton graphs is 2–5% faster than FA2 graphs here, so the gain +is graph dispatch reduction rather than the FA2 kernel; async+graphs is 1–4% +faster than FA2 graphs. That is a launch-overhead result on one L20, not a +production throughput qualification. A GPU test keeps `prepare_step` and +`__call__` free of host synchronization across a plan-width bucket change (it +runs them under `torch.cuda.set_sync_debug_mode`), and a CPU test bounds eager +plan creation over a 1K-token decode. + +The stored GSM8K-87 baseline fingerprints `torch==2.13.0`; a run on another torch +version cannot reuse its 59/87 baseline. A 10-question custom subset on torch 2.8 +was used only as an early smoke (3/10 with both backends, 8/10 identical). + + The [full 1,319-question comparison](https://github.com/InferMatrixLab/vllm-rlt/blob/d2f7db20593cccfe666b3c04c3ee50abc8c39af7/docs/benchmarks/gsm8k-bf16-20260912.md) recorded **62.02% native versus 61.64% Transformers** with strict matching. It includes per-question audits and a separate extraction diagnostic. That experiment's diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md new file mode 100644 index 00000000..19eee527 --- /dev/null +++ b/docs/flashinfer-fa2.md @@ -0,0 +1,193 @@ +# FlashInfer FA2 backend (Ada / SM89) + +`--attention-backend flashinfer_fa2` runs Ouro's paged decode through +FlashInfer's **FA2** paged-decode kernel. It targets NVIDIA Ada (SM89, +L20/L40S/RTX 40), where FlashAttention-2 rejects the page size 16 used by +Ouro's cache and FlashInfer XQA (PR #112) rejects SM89. + +## Scope + +- Ouro-1.4B BF16, fixed four-loop, decode-only. +- `block_size=16`, `head_dim=128`, NHD KV layout, BF16. +- LAST_EXITED and SHARED KV layouts. +- Eager execution, CUDA Graphs, and asynchronous scheduling. +- No packed prefill, prefill/decode disaggregation or speculative paths. +- Explicit opt-in only; `auto` keeps its existing FA4/FA3/FA2/Triton/Torch order. + +## Requirements + +- CUDA toolkit with `nvcc >= 12.9` reachable through `CUDA_HOME`, and that + toolkit's `lib` directory on `LD_LIBRARY_PATH` (FlashInfer JIT). +- `flashinfer-python>=0.7.2,<0.8`: + +```bash +pip install -e ".[flashinfer-fa2]" +``` + +The first call compiles the FA2 module for SM89 and caches it under +`~/.cache/flashinfer//89/`. + +## Usage + +```python +with LLM( + model_path, + device="cuda", + dtype=torch.bfloat16, + attention_backend="flashinfer_fa2", + cache_config=CacheConfig(num_blocks=512, block_size=16, layout="last_exited"), + scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=128), + execution_config=ExecutionConfig(cuda_graphs=True, async_scheduling=True), +) as llm: + ... +``` + +or `--attention-backend flashinfer_fa2` on the offline CLI/server, combined with +the engine's ordinary CUDA Graph and asynchronous-scheduling options. + +## Design + +- **Tensor-core, graph-compatible wrapper.** The backend uses + `CUDAGraphBatchDecodeWithPagedKVCacheWrapper(... use_tensor_cores=True)` for + both eager and graph execution. The CUDA-core decode wrapper bakes its work + schedule (request/tile indices) into plan-time state; on SM89 a plan built for + full capacity and replayed with a short live context stalls indefinitely once + the planned table is wide enough (reproduced at 64+ pages, GQA and MHA). The + tensor-core path plans against the planned width but takes the schedule from + the live CSR buffers. +- **Split-KV is disabled.** With FlashInfer's default splitting, a row's + reduction order depends on the planned width and on batch composition + (verified on L20: the same 1,500-token row differs between a 128- and a + 256-page plan, and alone versus inside a batch of eight). `disable_split_kv` + makes the schedule per-row: eager, graph replay, plan width and batch + composition are bit-identical. The cost is low-batch long-context throughput + (fewer CTAs); measured kernel-only at batch 8 it is unchanged. +- **Plan cache.** One plan per `(rows, q_heads, kv_heads, width bucket)`, where + the width is rounded up to a power of two. A growing eager context therefore + crosses at most one plan per bucket (~7 plans over a 1K-token request) instead + of one per page boundary. Each plan owns its CSR buffers and a workspace sized + by `workspace_size(...)`; `plan()`/`workspace_size()` inputs are pinned host + tensors, so creating a plan does not copy device metadata back to the host. + Plans captured by a CUDA Graph are pinned; eager plans are evicted first when + the cache (32 entries) is full. `plan_creations` reports plan count next to + `plan_misses`. +- **Kernel pinned by assertion.** The wrapper is constructed with + `backend="auto"` (the graph wrapper exposes no backend argument); after the + first `plan()` the adapter requires `resolved_backend == "fa2"` and records it + in `info`, so a patch release that routes to cuDNN fails loudly instead of + changing arithmetic. +- **Per-traversal metadata.** `prepare_attention` / `ModelRuntime.forward` prime + the backend once per traversal; the graph paths prime the *graph entry's* + private metadata after copying the caller's tables in. CSR construction + (`indptr`, compact `indices`, `last_page_len`) is fixed-shape and device-only, + and a GPU regression test runs it under `torch.cuda.set_sync_debug_mode` so a + step with device-owned context lengths cannot regain a host read. +- **Replay refreshes the plan buffers.** The first call of a traversal copies the + staging CSR into the plan's buffers and later layers skip it; a capture + records one copy group, so every replay consumes the metadata written by + `prepare_step` immediately before it. The staging buffers are referenced + weakly by their metadata tensors and stay alive as long as a graph entry does. +- **Padding.** Rows with `context_lengths == 0` are masked to exact zeros after + the kernel call. +- K/V must have identical strides. The engine's `key_cache[:, layer]` and + `value_cache[:, layer]` do; a contiguous K with a strided V is rejected by + FlashInfer. +- No `prefill` entry point: the default Ouro path never calls it, and + `packed_prefill` stays false. Asynchronous *speculative* scheduling requires + `device_context_lengths`, which this backend does not declare and is + therefore rejected at selection. + +## Validation (L20, 2026-11-01/02) + +Environment: 1×L20 (sm_89, 46 GB, driver 580.126.20), **torch 2.13.0** +(metadata `2.13.0`, CUDA 13.0), FlashInfer 0.7.2.post1, Triton 3.7.1, and the +protocol's pinned `transformers 4.55.0` / `lm-eval 0.4.9.2` / `datasets 3.6.0` +/ `tokenizers 0.21.4`. FlashInfer's JIT requires a CUDA >= 12.9 toolkit; the +official run used the pip `cuda-toolkit[nvcc]==13.0.3` tree. + +- `tests/test_flashinfer_fa2_attention.py` (22 cases: 16 GPU, 6 CPU): selection + and rejection paths; strided/GQA/MHA KV against both Torch and Triton kernels; + padding tri-state; zero-length rows; batch composition; vectorized CSR + against a reference construction (including the empty-row last-page length); + no host synchronization in `prepare_step`/`__call__` across a width-bucket + change; bounded plan creation over a 1K-token eager decode; bitwise + width/batch invariance and graph-replay equality at a 128-page plan; and + tiny-Ouro runs where eager, graphs, async and async+graphs produce identical + token/exit-depth sequences for LAST_EXITED and SHARED. +- Real Ouro-1.4B BF16, greedy, trace depth 4, three requests: Triton eager, + FA2 eager, FA2 graphs, and FA2 graphs+async produce **identical token + sequences**; `plan_misses == 0`, one graph capture, 60 replays. +- Official GSM8K-87 (default protocol, HF baseline 59/87): + +| run | correct | accuracy | gate | +|---|---:|---:|---| +| Triton eager | 61/87 | 70.11% | pass | +| FA2 eager (tensor cores) | 59/87 | 67.82% | pass (+0.00 pp) | +| FA2 + CUDA Graphs | 59/87 | 67.82% | pass (+0.00 pp) | + + Eager and graphs are **bit-identical: 87/87 token sequences**, with no + length-limited runaway (mean 106.7 tokens; eager 717.8 s, graphs 279.6 s). + Against Triton, 67/87 sequences are identical; four answers flip (203, 426, + 819 correct → incorrect; 871 incorrect → correct). The no-split arithmetic + reproduces the earlier split-KV *graph* run on all 87 questions; the earlier + split-KV eager run differs on 18 questions because its reduction order + depended on the narrow eager plan width. Disabling split-KV moved the eager + score from 60/87 to 59/87, exactly at the baseline. +- The earlier CUDA-core variant of this branch scored 61/87 with a different + BF16 accumulation order; it is superseded by the tensor-core kernel that + graph execution requires. +- An earlier 10-question custom subset on torch 2.8 was used as a smoke: + 3/10 with both backends and 8/10 identical sequences. + +## Performance (diagnostic, not a benchmark) + +Matched fixed work, real Ouro-1.4B BF16, greedy depth 4, 64 output tokens, +32-token prompts, one wave of N requests, second run after a warmup that +captures the graph. All five columns come from one session: + +| concurrency | Triton eager | Triton graphs | FA2 eager | FA2 graphs | FA2 async+graphs | +|---:|---:|---:|---:|---:|---:| +| 1 | 5.20 s (12.3 tok/s) | 1.52 s (42.2) | 5.49 s (11.7) | 1.59 s (40.2) | 1.58 s (40.6) | +| 8 | 5.34 s (95.8) | 1.76 s (290.2) | 5.55 s (92.2) | 1.84 s (278.1) | 1.81 s (282.8) | +| 32 | 6.26 s (327.3) | 2.20 s (930.1) | 6.32 s (324.1) | 2.24 s (915.4) | 2.15 s (952.3) | + +Eager FA2 is at parity with Triton eager. CUDA Graphs remove most per-step +launch overhead for **both** backends (2.8–3.5× over eager on this small +model); Triton graphs is then 2–5% faster than FA2 graphs, so the graph gain is +dispatch reduction, not the FA2 kernel. Asynchronous scheduling adds 1–4% over +FA2 graphs at these batch sizes. The official GSM8K-87 runs agree in direction +(eager 717.8 s, graphs 279.6 s, identical generated lengths). These are +single-GPU, single-wave diagnostics, not production throughput claims. + +Disabling split-KV costs attention time only at low batch and long context +(kernel-only, one L20, page 16, random KV): + +| rows | ctx | split-KV ms | no split ms | ratio | +|---:|---:|---:|---:|---:| +| 1 | 4096 | 0.023 | 0.049 | 2.13× | +| 1 | 8192 | 0.028 | 0.092 | 3.35× | +| 8 | 4096 | 0.375 | 0.372 | 0.99× | +| 8 | 8192 | 0.741 | 0.739 | 1.00× | + +Ouro runs 24 attention layers per depth view, so at batch 1 and 8K context the +no-split schedule adds roughly 6 ms per generated token; at batch 8 it is +unchanged. The determinism above is the reason for accepting that cost. + +Kernel-only figures from the earlier CUDA-core adapter (FA2/Triton 0.39–0.98 +across B 1–32 and context 256–8192) are superseded by the tensor-core kernel +used here and are not reproduced. + +## Follow-ups + +- Auto-selection on SM89 after broader validation. +- More page sizes/head dimensions, packed prefill, PD and speculative paths. +- Revisit the CUDA-core decode path if FlashInfer makes its plan schedule + independent of plan-time lengths; it is bit-closer to the Torch reference. +- Write the live CSR directly into the plan buffers at prime time: one copy per + traversal remains because the plan key needs head counts that are only known + at the first call. +- Revisit split-KV for low-batch long-context workloads if FlashInfer offers a + deterministic split schedule; today the deterministic path costs 2.1–3.4× + kernel time there. +- Reconcile with #112 if XQA lands first; this backend keeps a distinct + `flashinfer_fa2` name and module. diff --git a/pyproject.toml b/pyproject.toml index 04ccfa6f..cd45f12f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ pd = ["nixl==1.4.1"] serve = [] flash-attn = ["flash-attn>=2.8.3"] flash-attn4 = ["flash-attn-4==4.0.0b30"] +flashinfer-fa2 = ["flashinfer-python>=0.7.2,<0.8"] triton = [] dev = ["pytest>=8", "ruff>=0.9"] diff --git a/tests/test_flashinfer_fa2_attention.py b/tests/test_flashinfer_fa2_attention.py new file mode 100644 index 00000000..113cbe8f --- /dev/null +++ b/tests/test_flashinfer_fa2_attention.py @@ -0,0 +1,472 @@ +"""FlashInfer FA2 adapter: SM89 selection, CSR metadata and eager decode.""" + +import dataclasses +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SamplingParams, + SchedulerConfig, +) +from vllm_rlt.attention import BackendCapabilities, create_backend, validate_backend_name +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.kernels import flashinfer_fa2_attention as fa2 +from vllm_rlt.kernels.flashinfer_fa2_attention import FlashInferFA2PagedAttention +from vllm_rlt.kernels.paged_attention import torch_paged_attention +from vllm_rlt.models import OuroConfig, OuroForCausalLM + + +@pytest.fixture +def fake_flashinfer(monkeypatch): + package = types.ModuleType("flashinfer") + package.__path__ = [] + decode = types.ModuleType("flashinfer.decode") + + class Wrapper: + def __init__(self, *args, **kwargs): + pass + + decode.BatchDecodeWithPagedKVCacheWrapper = Wrapper + decode.CUDAGraphBatchDecodeWithPagedKVCacheWrapper = Wrapper + package.decode = decode + monkeypatch.setitem(sys.modules, "flashinfer", package) + monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) + monkeypatch.setattr(fa2, "version", lambda name: "0.7.2.post1") + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) + + +def test_backend_name_and_explicit_selection(fake_flashinfer): + validate_backend_name("flashinfer_fa2") + selected = create_backend("flashinfer_fa2", torch.device("cuda"), torch.bfloat16, 128, 16) + assert selected.info["selected_backend"] == "flashinfer_fa2" + capabilities = selected.capabilities + assert capabilities.cuda_graphs + assert capabilities.async_scheduling + assert not capabilities.packed_prefill + assert not capabilities.device_context_lengths + assert selected.info["cuda_graphs"] and selected.info["async_scheduling"] + create_backend( + "flashinfer_fa2", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=BackendCapabilities(cuda_graphs=True, async_scheduling=True), + ) + with pytest.raises(ValueError, match="device_context_lengths"): + create_backend( + "flashinfer_fa2", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=BackendCapabilities(device_context_lengths=True), + ) + + +def test_constructor_rejects_unvalidated_shapes(fake_flashinfer, monkeypatch): + with pytest.raises(ValueError, match="SM89"): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 6)) + FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 16) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) + with pytest.raises(ValueError, match="bfloat16"): + FlashInferFA2PagedAttention(torch.device("cuda"), torch.float16, 128, 16) + with pytest.raises(ValueError, match="block-size"): + FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 32) + with pytest.raises(ValueError, match="head_dim"): + FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 64, 16) + + +def test_missing_package_never_falls_back(fake_flashinfer, monkeypatch): + def missing(name): + raise fa2.PackageNotFoundError(name) + + monkeypatch.setattr(fa2, "version", missing) + with pytest.raises(ImportError, match="flashinfer-python"): + create_backend("flashinfer_fa2", torch.device("cuda"), torch.bfloat16, 128, 16) + + +def test_call_rejects_row_mismatch(fake_flashinfer): + kernel = create_backend("flashinfer_fa2", torch.device("cuda"), torch.bfloat16, 128, 16).kernel + q = torch.empty(3, 4, 128, dtype=torch.bfloat16) + tables = torch.zeros(2, 4, dtype=torch.int32) + lengths = torch.zeros(2, dtype=torch.int32) + with pytest.raises(ValueError, match="one metadata row per query row"): + kernel(q, None, None, tables, lengths) + + +def test_wrapper_rejects_unsupported_head_ratio(fake_flashinfer): + kernel = create_backend("flashinfer_fa2", torch.device("cuda"), torch.bfloat16, 128, 16).kernel + with pytest.raises(ValueError, match="head ratio"): + kernel._wrapper(2, 3, 1, 4) + + +def test_plan_count_is_bounded(monkeypatch): + plans = {"count": 0} + + class Wrapper: + def __init__(self, *args, **kwargs): + pass + + def workspace_size(self, *args, **kwargs): + return 1024, 1024 + + def reset_workspace_buffer(self, *args): + pass + + def plan(self, *args, **kwargs): + plans["count"] += 1 + + def run(self, q, kv, out): + out.zero_() + + @property + def resolved_backend(self): + return "fa2" + + package = types.ModuleType("flashinfer") + package.__path__ = [] + decode = types.ModuleType("flashinfer.decode") + decode.CUDAGraphBatchDecodeWithPagedKVCacheWrapper = Wrapper + package.decode = decode + monkeypatch.setitem(sys.modules, "flashinfer", package) + monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) + monkeypatch.setattr(fa2, "version", lambda name: "0.7.2.post1") + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 16) + attention.device = torch.device("cpu") + + def step(lengths): + rows = len(lengths) + width = max((length - 1) // 16 + 1 for length in lengths) + tables = torch.randint(0, 1000, (rows, width), dtype=torch.int32) + context = torch.tensor(lengths, dtype=torch.int32) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=context),)) + q = torch.zeros(rows, 16, 128, dtype=torch.bfloat16) + kv = torch.zeros(1, 16, 16, 128, dtype=torch.bfloat16) + attention(q, kv, kv, tables, context) + + position = 0 + while position < 700: + count = min(128, 700 - position) + step(list(range(position + 1, position + count + 1))) + position += count + for index in range(300): + step([701 + index]) + + assert plans["count"] == attention.plan_creations + # Power-of-two width buckets: 1, 2, 4, 8, 16, 32, 64 for this run. + assert attention.plan_creations <= 10 + assert attention.plan_entries <= 10 + assert attention.plan_misses == 0 + assert attention.info["resolved_backend"] == "fa2" + + +@pytest.mark.gpu +def test_strided_gqa_padding_and_zero_rows(): + torch.manual_seed(7) + page, head_dim, q_heads, kv_heads = 16, 128, 4, 2 + # Exactly KVCacheManager.key_cache[:, layer]: non-contiguous block stride. + keys = torch.randn(12, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + values = torch.randn(12, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + q = torch.randn(4, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + tables = torch.tensor( + [[7, 2, 9], [3, 0, 0], [1, 8, 6], [0, 0, 0]], dtype=torch.int32, device="cuda" + ) + lengths = torch.tensor([35, 1, 47, 0], dtype=torch.int32, device="cuda") + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=lengths),)) + actual = attention(q, keys, values, tables, lengths) + expected = torch_paged_attention(q, keys, values, tables, lengths) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=0) + assert torch.isfinite(actual).all() + assert torch.count_nonzero(actual[-1]) == 0 + assert attention.plan_misses == 0 + assert attention.plan_entries == 1 + + for padding in (-1, 0, 9): + padded = tables.clone() + padded[1, 1:] = padding + padded[3, :] = padding + attention.prepare_step((SimpleNamespace(block_tables=padded, context_lengths=lengths),)) + out = attention(q, keys, values, padded, lengths) + assert torch.equal(out[:3], actual[:3]) + assert torch.count_nonzero(out[3]) == 0 + + subset_tables, subset_lengths = tables[:3], lengths[:3] + attention.prepare_step( + (SimpleNamespace(block_tables=subset_tables, context_lengths=subset_lengths),) + ) + smaller = attention(q[:3], keys, values, subset_tables, subset_lengths) + assert torch.equal(smaller, actual[:3]) + assert attention.plan_misses == 0 + + +@pytest.mark.gpu +def test_empty_batch_and_unprimed_fallback(): + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 16) + keys = torch.randn(4, 16, 2, 128, device="cuda", dtype=torch.bfloat16) + empty = attention( + torch.empty(0, 4, 128, device="cuda", dtype=torch.bfloat16), + keys, + keys, + torch.empty(0, 1, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + ) + assert empty.shape == (0, 4, 128) + + q = torch.randn(1, 4, 128, device="cuda", dtype=torch.bfloat16) + tables = torch.tensor([[0, 1]], dtype=torch.int32, device="cuda") + lengths = torch.tensor([32], dtype=torch.int32, device="cuda") + expected = torch_paged_attention(q, keys, keys, tables, lengths) + actual = attention(q, keys, keys, tables, lengths) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=0) + assert attention.plan_misses == 1 + + +def _tiny_config(): + return OuroConfig( + vocab_size=64, + hidden_size=512, + intermediate_size=1024, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=128, + max_position_embeddings=128, + total_ut_steps=4, + bos_token_id=0, + eos_token_id=0, + ) + + +@pytest.mark.gpu +def test_vectorized_csr_matches_reference(): + torch.manual_seed(11) + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 16) + tables = torch.randint(0, 40, (6, 10), device="cuda", dtype=torch.int32) + lengths = torch.tensor([1, 15, 16, 17, 129, 0], device="cuda", dtype=torch.int32) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=lengths),)) + csr = attention._csr[(id(tables), id(lengths))] + pages = ((lengths + 15) // 16).tolist() + rows = [tables[row, :count].tolist() for row, count in enumerate(pages)] + assert csr.indptr.tolist() == [0] + torch.cumsum(torch.tensor(pages, device="cuda"), 0).tolist() + assert csr.indices[: csr.indptr[-1]].tolist() == [page for row in rows for page in row] + # Empty rows keep the kernel's required last-page length of one. + assert csr.last_page_len.tolist() == [ + 1 if count == 0 else int(length) - (count - 1) * 16 + for length, count in zip(lengths.tolist(), pages) + ] + + +@pytest.mark.gpu +def test_prepare_step_does_not_synchronize(): + page, head_dim, rows = 16, 128, 2 + torch.manual_seed(17) + keys = torch.randn(64, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + values = torch.randn(64, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + q = torch.randn(rows, 16, head_dim, device="cuda", dtype=torch.bfloat16) + lengths = torch.tensor([33, 47], device="cuda", dtype=torch.int32) + # 64 and 65 pages land in adjacent power-of-two plan buckets. + batches = [] + for width in (64, 65): + tables = torch.randint(0, 64, (rows, width), device="cuda", dtype=torch.int32) + batches.append(SimpleNamespace(block_tables=tables, context_lengths=lengths)) + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + # Warm both plans, module load and allocator caches outside the measured region. + for batch in batches: + attention.prepare_step((batch,)) + attention(q, keys, values, batch.block_tables, lengths) + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + for batch in batches: + attention.prepare_step((batch,)) + attention(q, keys, values, batch.block_tables, lengths) + finally: + torch.cuda.set_sync_debug_mode("default") + assert attention.plan_creations == 2 + + +@pytest.mark.gpu +def test_graph_replay_reads_refreshed_metadata(): + torch.manual_seed(13) + page, head_dim, q_heads, kv_heads = 16, 128, 4, 2 + keys = torch.randn(8, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + values = torch.randn(8, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + q = torch.randn(2, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + tables = torch.zeros((2, 4), dtype=torch.int32, device="cuda") + lengths = torch.tensor([16, 32], dtype=torch.int32, device="cuda") + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=lengths),)) + attention(q, keys, values, tables, lengths) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = attention(q, keys, values, tables, lengths) + tables.copy_(torch.tensor([[7, 2, 3, 0], [3, 0, 0, 0]], dtype=torch.int32, device="cuda")) + lengths.copy_(torch.tensor([35, 1], dtype=torch.int32, device="cuda")) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=lengths),)) + graph.replay() + torch.cuda.synchronize() + expected = torch_paged_attention(q, keys, values, tables, lengths) + torch.testing.assert_close(output, expected, atol=1e-2, rtol=0) + assert attention.plan_misses == 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize(("q_heads", "kv_heads"), [(4, 2), (16, 16)]) +def test_paged_attention_matches_triton(q_heads, kv_heads): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + torch.manual_seed(7) + page, head_dim, rows = 16, 128, 5 + # Strided views exactly like KVCacheManager.key_cache[:, layer]. + keys = torch.randn(12, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + values = torch.randn(12, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + q = torch.randn(rows, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + tables = torch.tensor( + [[7, 2, 9, 3], [3, 0, 0, 0], [1, 8, 6, 4], [0, 0, 0, 0], [5, 5, 5, 5]], + dtype=torch.int32, + device="cuda", + ) + lengths = torch.tensor([35, 1, 47, 0, 64], dtype=torch.int32, device="cuda") + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + attention.prepare_step((SimpleNamespace(block_tables=tables, context_lengths=lengths),)) + actual = attention(q, keys, values, tables, lengths) + expected = TritonPagedAttention()(q, keys, values, tables, lengths) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=0) + torch.testing.assert_close( + actual, torch_paged_attention(q, keys, values, tables, lengths), atol=1e-2, rtol=0 + ) + assert torch.count_nonzero(actual[3]) == 0 + assert attention.plan_misses == 0 + + +@pytest.mark.gpu +def test_attention_is_width_and_batch_invariant(): + """A row's output must not depend on plan width, batch size or graph replay.""" + page, head_dim, q_heads, kv_heads, width, blocks = 16, 128, 16, 16, 128, 512 + torch.manual_seed(21) + keys = torch.randn(blocks, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[ + :, 1 + ] + values = torch.randn(blocks, 3, page, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)[ + :, 1 + ] + q0 = torch.randn(1, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + q8 = torch.cat( + [q0] + + [torch.randn(1, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) for _ in range(7)] + ) + tables0 = torch.randint(0, blocks, (1, width), device="cuda", dtype=torch.int32) + tables8 = torch.cat( + [tables0] + + [torch.randint(0, blocks, (1, width), device="cuda", dtype=torch.int32) for _ in range(7)] + ) + lengths0 = torch.tensor([1500], device="cuda", dtype=torch.int32) + lengths8 = torch.tensor( + [1500, 900, 1200, 700, 1100, 800, 1000, 600], device="cuda", dtype=torch.int32 + ) + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + attention.prepare_step((SimpleNamespace(block_tables=tables0, context_lengths=lengths0),)) + out0 = attention(q0, keys, values, tables0, lengths0) + attention.prepare_step((SimpleNamespace(block_tables=tables8, context_lengths=lengths8),)) + out8 = attention(q8, keys, values, tables8, lengths8) + assert torch.equal(out0[0], out8[0]) + assert attention.plan_creations == 2 + + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = attention(q8, keys, values, tables8, lengths8) + new_tables = torch.randint(0, blocks, (8, width), device="cuda", dtype=torch.int32) + new_tables[0] = tables0[0] + new_lengths = torch.tensor( + [1400, 900, 1200, 700, 1100, 800, 1000, 600], device="cuda", dtype=torch.int32 + ) + tables8.copy_(new_tables) + lengths8.copy_(new_lengths) + attention.prepare_step((SimpleNamespace(block_tables=tables8, context_lengths=lengths8),)) + graph.replay() + eager = attention(q8, keys, values, tables8, lengths8) + torch.cuda.synchronize() + assert torch.equal(captured, eager) + + +def _tiny_ouro_outputs(model, layout, execution): + traces = {str(i): [4, 2 + i % 3, 4, 2, 3, 4] for i in range(3)} + engine = StageEngine( + model, + cache_config=CacheConfig(num_blocks=64, block_size=16, layout=layout), + scheduler_config=SchedulerConfig( + max_num_seqs=3, max_num_batched_tokens=8, prefill_chunk_size=4 + ), + execution_config=execution, + exit_config=ExitConfig("trace", depths_by_request=traces), + attention_backend="flashinfer_fa2", + ) + for i in range(3): + engine.add_request( + str(i), + [2, 3, 4, 5] * (i + 1), + SamplingParams(max_tokens=6, min_loops=1, ignore_eos=True), + ) + finished = {} + for _ in range(300): + if not engine.has_unfinished_requests(): + break + for out in engine.step(): + if out.finished: + finished[out.request_id] = (out.token_ids, out.exit_depths) + assert len(finished) == 3 + assert engine.cache_manager.num_used_blocks == 0 + assert engine.cache_manager.attention.kernel.plan_misses == 0 + return engine, finished + + +@pytest.mark.gpu +@pytest.mark.parametrize("layout", ["last_exited", "shared"]) +@pytest.mark.parametrize( + "execution", + [ + ExecutionConfig(cuda_graphs=True), + ExecutionConfig(async_scheduling=True), + ExecutionConfig(async_scheduling=True, cuda_graphs=True), + ], + ids=["graphs", "async", "async+graphs"], +) +def test_tiny_ouro_execution_variants_match_eager(layout, execution): + """The TF2 tensor-core kernel is the shared arithmetic; only routing changes.""" + torch.manual_seed(123) + model = OuroForCausalLM(_tiny_config()).to(device="cuda", dtype=torch.bfloat16) + reference, expected = _tiny_ouro_outputs(model, layout, ExecutionConfig()) + engine, actual = _tiny_ouro_outputs(model, layout, execution) + assert actual == expected + if execution.cuda_graphs: + assert engine.model_runner.graphs.captures > 0 + assert engine.model_runner.graphs.replays > engine.model_runner.graphs.captures + + +@pytest.mark.gpu +@pytest.mark.parametrize("layout", ["last_exited", "shared"]) +def test_tiny_ouro_graphs_match_eager_with_wide_tables(layout): + """Graph plans use the full 128-page table; eager plans use narrow buckets.""" + torch.manual_seed(123) + config = dataclasses.replace(_tiny_config(), max_position_embeddings=2048) + model = OuroForCausalLM(config).to(device="cuda", dtype=torch.bfloat16) + reference, expected = _tiny_ouro_outputs(model, layout, ExecutionConfig()) + engine, actual = _tiny_ouro_outputs(model, layout, ExecutionConfig(cuda_graphs=True)) + assert actual == expected + assert engine.model_runner.graphs.captures > 0 diff --git a/tests/test_gsm8k_accuracy.py b/tests/test_gsm8k_accuracy.py index c7da98b0..5472baf1 100644 --- a/tests/test_gsm8k_accuracy.py +++ b/tests/test_gsm8k_accuracy.py @@ -703,3 +703,15 @@ def test_compare_reports_sequences_without_changing_the_gate(tmp_path): assert differences["exit_depths"]["questions_differing_on_identical_context"] == 1 assert result["answer_disagreements"] == [] assert result["passes_observed_accuracy_gate"] + + +def test_cuda_graphs_requires_native_before_creating_output(tmp_path): + args = SimpleNamespace( + backend="transformers", + cuda_graphs=True, + protocol=str(tmp_path / "protocol.json"), + output=str(tmp_path / "out"), + ) + with pytest.raises(ValueError, match="native backend"): + gsm8k.run(args) + assert not (tmp_path / "out").exists() diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..a7c3147b 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -39,6 +39,12 @@ def __call__(self, *args, **kwargs): def prefill(self, *args, **kwargs): return self.kernel.prefill(*args, **kwargs) + def prepare_step(self, batches): + """Prime a stateful backend once per traversal; stateless kernels no-op.""" + prepare = getattr(self.kernel, "prepare_step", None) + if prepare is not None: + prepare(batches) + @property def generation(self): # Retain the existing Flash adapter's diagnostic attribute. @@ -60,6 +66,7 @@ def generation(self): device_context_lengths=True, ), } +_FLASHINFER_FA2_CAPABILITIES = BackendCapabilities(cuda_graphs=True, async_scheduling=True) _AUTO_PRIORITY = ("flash_attn_4", "flash_attn_3", "flash_attn_2", "triton", "torch") @@ -68,7 +75,7 @@ def backend_capabilities(implementation) -> BackendCapabilities: def validate_backend_name(backend: str) -> None: - if backend not in {"auto", "torch", "triton", *FLASH_BACKENDS}: + if backend not in {"auto", "torch", "triton", "flashinfer_fa2", *FLASH_BACKENDS}: raise ValueError("unknown attention backend") @@ -83,6 +90,13 @@ def _create_explicit(backend, device, dtype, head_dim, block_size): # Probe the actual module, including its dependency imports, before selection. importlib.import_module("vllm_rlt.kernels.triton_attention") return AttentionBackend(TritonPagedAttention(), _TRITON_CAPABILITIES, {"backend": backend}) + if backend == "flashinfer_fa2": + from vllm_rlt.kernels.flashinfer_fa2_attention import FlashInferFA2PagedAttention + + kernel = FlashInferFA2PagedAttention(device, dtype, head_dim, block_size) + # Share the dict so the resolved kernel recorded after the first plan is + # visible through the backend's selection info. + return AttentionBackend(kernel, _FLASHINFER_FA2_CAPABILITIES, kernel.info) if backend in FLASH_BACKENDS: kernel = FlashPagedAttention(device, dtype, head_dim, block_size, backend) return AttentionBackend(kernel, _FLASH_CAPABILITIES[kernel.generation], dict(kernel.info)) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..b8158f62 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -628,6 +628,10 @@ def _require_live_batch(self, batch: KVMetadata | DeviceKVMetadata) -> None: if self._allocations.get(request_id) is not allocation: raise RuntimeError(f"stale prepared KV batch for request {request_id!r}") + def prepare_attention(self, batch: KVMetadata | DeviceKVMetadata) -> None: + """Prime a stateful attention backend once for this traversal.""" + self.attention.prepare_step((batch,)) + def _validate_tensor(self, tensor, batch_size, name, *, query=False): if tensor.ndim != 3 or tensor.shape[0] != batch_size or tensor.shape[-1] != self.head_dim: raise ValueError(f"{name} must have shape [batch, heads, {self.head_dim}]") @@ -709,6 +713,7 @@ def attend( """Apply causal GQA at each row's own position and recurrence depth.""" self._validate_layer(layer) batch = self._prepare_batch(request_ids, depths, positions, for_write=False) + self.prepare_attention(batch) return self._attend_prepared(layer, batch, q) @torch.no_grad() diff --git a/vllm_rlt/entrypoints/cli.py b/vllm_rlt/entrypoints/cli.py index a9b8025c..7ebb5f53 100644 --- a/vllm_rlt/entrypoints/cli.py +++ b/vllm_rlt/entrypoints/cli.py @@ -74,6 +74,7 @@ def main(): "flash_attn_2", "flash_attn_3", "flash_attn_4", + "flashinfer_fa2", ], default="torch", ) diff --git a/vllm_rlt/entrypoints/serve.py b/vllm_rlt/entrypoints/serve.py index cd452d92..18955b80 100644 --- a/vllm_rlt/entrypoints/serve.py +++ b/vllm_rlt/entrypoints/serve.py @@ -95,6 +95,7 @@ def main(): "flash_attn_2", "flash_attn_3", "flash_attn_4", + "flashinfer_fa2", ], default="triton", ) diff --git a/vllm_rlt/kernels/flashinfer_fa2_attention.py b/vllm_rlt/kernels/flashinfer_fa2_attention.py new file mode 100644 index 00000000..4269c5c6 --- /dev/null +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -0,0 +1,382 @@ +"""FlashInfer FA2 paged decode for Ada (SM89); eager, CUDA Graphs and async steps. + +The adapter builds the CSR metadata FlashInfer's paged kernel wants once per +traversal (``prepare_step``), then reuses it across every decoder layer through +``__call__``. Plans are keyed by ``(rows, heads, width bucket)``, where the +width is rounded up to a power of two, so a growing eager context creates at +most one plan per bucket instead of one per page boundary. + +Split-KV is disabled: with it, a row's reduction order depends on the planned +width and on batch composition, so eager and graph execution (padded rows, +full-width graph plans) disagree bit for bit. Without splits the schedule is +per-row, and eager, graphs and any batch composition produce identical output. +The cost is low-batch long-context throughput (fewer CTAs); at batch 8 the +kernel cost is unchanged. + +Only FlashInfer's CUDA-Graph-compatible tensor-core FA2 wrapper is used. The +CUDA-core decode wrapper bakes its work schedule (request/tile indices) into +plan-time state, so a plan built for full capacity and replayed with a short +live context stalls on SM89; the tensor-core path takes the schedule from the +live CSR buffers and masks unused splits, which is exactly the graph contract. +The cost is FlashInfer's documented caveat that it does not switch kernels per +batch size or sequence length. + +Construction is device-only: the per-row page counts, the compact page index +list and the last-page lengths are all derived with fixed-shape kernels, so a +step with device-owned context lengths never reads back to the host. + +The CSR staging buffers stay allocated for as long as their metadata tensors +are alive (weak references). CUDA Graph capture records the copy from those +buffers into the wrapper's plan buffers, so replay refreshes the plan buffers +from whatever ``prepare_step`` wrote before the replay; no Python runs inside +the graph. + +Only the L20/L40S/RTX-40 class is accepted: FlashAttention-2 rejects the page +size 16 used by Ouro's cache, and FlashInfer XQA (PR #112) rejects SM89. +""" + +import weakref +from importlib.metadata import PackageNotFoundError, version + +import torch + +FLASHINFER_FA2_PAGE_SIZE = 16 +FLASHINFER_FA2_HEAD_DIM = 128 +FLASHINFER_FA2_GROUP_SIZES = (1, 2, 4, 8) +# Live metadata sets are tiny (graph entries plus the current eager step). The +# bound only guards a caller that keeps many distinct tensors alive. +FLASHINFER_FA2_CSR_CAPACITY = 64 +# Each (rows, heads, width) plan owns a device workspace. Reuse across the +# engine's batch sizes; captured plans are pinned and never evicted. +FLASHINFER_FA2_PLAN_CAPACITY = 32 + + +def _version_tuple(value): + core = value.split("+", 1)[0].split(".post", 1)[0] + parts = [] + for component in core.split(".")[:3]: + digits = "".join(character for character in component if character.isdigit()) + parts.append(int(digits) if digits else 0) + return tuple(parts) + + +class _CSR: + """CSR sources owned by one (block_tables, context_lengths) pair. + + Metadata is referenced weakly so an eager step can release it; a graph entry + keeps its metadata alive for the runner's lifetime, which keeps these sources + and the raw pointers recorded by its capture valid. + """ + + __slots__ = ( + "indptr", + "indices", + "last_page_len", + "block_tables", + "context_lengths", + "synced", + "captured", + ) + + def __init__(self, block_tables, context_lengths): + rows, width = block_tables.shape + device = block_tables.device + self.indptr = torch.empty(rows + 1, dtype=torch.int32, device=device) + self.indices = torch.empty(rows * width, dtype=torch.int32, device=device) + self.last_page_len = torch.empty(rows, dtype=torch.int32, device=device) + self.block_tables = weakref.ref(block_tables) + self.context_lengths = weakref.ref(context_lengths) + # Plans already refreshed from these sources in the current traversal + # (eager) or recorded into a capture, so layers do not copy again. + self.synced = set() + self.captured = set() + + def owns(self, block_tables, context_lengths): + return self.block_tables() is block_tables and self.context_lengths() is context_lengths + + +class _Plan: + """One planned wrapper with the buffers its captured replays read. + + The wrapper keeps the CSR buffers by reference; ``prepare_step`` refreshes + the staging copy and the first call of a traversal copies it in, so later + layers and replays see the same metadata without per-layer copies. + """ + + __slots__ = ("wrapper", "indptr", "indices", "last_page_len", "pinned") + + def __init__(self, wrapper, indptr, indices, last_page_len, pinned): + self.wrapper = wrapper + self.indptr = indptr + self.indices = indices + self.last_page_len = last_page_len + self.pinned = pinned + + +class FlashInferFA2PagedAttention: + def __init__(self, device, dtype, head_dim, block_size): + if device.type != "cuda" or torch.version.hip: + raise ValueError("flashinfer_fa2 requires an NVIDIA CUDA device") + if dtype is not torch.bfloat16: + raise ValueError("flashinfer_fa2 requires bfloat16") + if block_size != FLASHINFER_FA2_PAGE_SIZE: + raise ValueError(f"flashinfer_fa2 requires --block-size {FLASHINFER_FA2_PAGE_SIZE}") + if head_dim != FLASHINFER_FA2_HEAD_DIM: + raise ValueError(f"flashinfer_fa2 requires head_dim {FLASHINFER_FA2_HEAD_DIM}") + capability = torch.cuda.get_device_capability(device) + if capability != (8, 9): + raise ValueError( + f"flashinfer_fa2 is validated on SM89 only, got SM{capability[0]}{capability[1]}" + ) + try: + package_version = version("flashinfer-python") + if not (0, 7, 2) <= _version_tuple(package_version) < (0, 8): + raise ImportError( + "flashinfer_fa2 requires flashinfer-python>=0.7.2,<0.8; " + f"found {package_version}" + ) + from flashinfer.decode import CUDAGraphBatchDecodeWithPagedKVCacheWrapper + except (ImportError, PackageNotFoundError) as error: + raise ImportError( + "install flashinfer with `pip install 'flashinfer-python>=0.7.2,<0.8'` " + "(CUDA >= 12.9 is required for its JIT kernels); no backend fallback " + "was performed" + ) from error + self._wrapper_cls = CUDAGraphBatchDecodeWithPagedKVCacheWrapper + self.device = torch.device(device) + self.head_dim = head_dim + self.page_size = block_size + self._wrappers = {} + self._column_index = {} + self._csr = {} + self._primed = {} + self._plan_misses = 0 + self._plan_creations = 0 + self.info = dict( + backend="flashinfer_fa2", + package="flashinfer-python", + version=package_version, + kernel="fa2", + tensor_cores=True, + cuda_graph_wrapper=True, + split_kv=False, + page_size=block_size, + head_dim=head_dim, + sm="8.9", + cuda_graphs=True, + async_scheduling=True, + ) + + @property + def plan_misses(self): + return self._plan_misses + + @property + def plan_creations(self): + return self._plan_creations + + @property + def plan_entries(self): + return len(self._wrappers) + + def prepare_step(self, batches): + """Build this traversal's CSR once per batch, without host synchronization. + + The graph paths call this with the graph entry's private metadata after + copying the caller's tables in, so eager and graph execution share one + contract. + """ + self._primed.clear() + for batch in batches: + block_tables = getattr(batch, "block_tables", None) + context_lengths = getattr(batch, "context_lengths", None) + if block_tables is None or context_lengths is None: + continue + self._validate_metadata(block_tables, context_lengths) + if block_tables.shape[0] == 0: + continue + csr = self._csr_entry(block_tables, context_lengths) + self._write_csr(csr, block_tables, context_lengths) + csr.synced.clear() + self._primed[(id(block_tables), id(context_lengths))] = csr + + @staticmethod + def _validate_metadata(block_tables, context_lengths): + if block_tables.ndim != 2 or context_lengths.ndim != 1: + raise ValueError("flashinfer_fa2 expects block_tables [rows, width] and lengths [rows]") + if block_tables.shape[0] != context_lengths.shape[0]: + raise ValueError("flashinfer_fa2 block_tables and context_lengths disagree on rows") + if block_tables.dtype != torch.int32 or context_lengths.dtype != torch.int32: + raise ValueError("flashinfer_fa2 expects int32 block_tables and context_lengths") + + def _csr_entry(self, block_tables, context_lengths): + key = (id(block_tables), id(context_lengths)) + csr = self._csr.get(key) + if csr is not None and csr.owns(block_tables, context_lengths): + return csr + csr = _CSR(block_tables, context_lengths) + self._csr[key] = csr + if len(self._csr) > FLASHINFER_FA2_CSR_CAPACITY: + for stale, entry in list(self._csr.items()): + if entry.block_tables() is None or entry.context_lengths() is None: + del self._csr[stale] + return csr + + def _columns(self, device, width): + key = (device, width) + columns = self._column_index.get(key) + if columns is None: + columns = torch.arange(width, dtype=torch.int32, device=device).unsqueeze(0) + self._column_index[key] = columns + return columns + + def _write_csr(self, csr, block_tables, context_lengths): + """Refresh CSR sources in place; every shape below is static. + + Valid pages of row ``r`` land in ``[indptr[r], indptr[r + 1])``; unused + slots receive padding columns, which the kernel never reads because it + stops at ``last_page_len``. + """ + width = block_tables.shape[1] + page = self.page_size + pages = (context_lengths.clamp_min(0) + page - 1) // page + # A Python scalar assignment to a CUDA slice stages through host memory + # and synchronizes; fill_ keeps CSR construction device-only. + csr.indptr[0].fill_(0) + csr.indptr[1:].copy_(torch.cumsum(pages, dim=0)) + valid = self._columns(block_tables.device, width) < pages.unsqueeze(1) + flat = valid.reshape(-1) + packed = torch.cumsum(flat.to(torch.int32), dim=0) - 1 + tail = flat.sum(dtype=torch.int32) + torch.cumsum((~flat).to(torch.int32), dim=0) - 1 + order = torch.where(flat, packed, tail).to(torch.int64) + csr.indices.scatter_(0, order, block_tables.reshape(-1)) + torch.where( + pages == 0, + torch.ones_like(pages), + context_lengths - (pages - 1) * page, + out=csr.last_page_len, + ) + + def _wrapper(self, rows, q_heads, kv_heads, width): + # Round the table width up to a power of two: a growing eager context + # then crosses at most one plan per bucket instead of one per page. + width = max(1, width) + width = 1 << (width - 1).bit_length() + key = (rows, q_heads, kv_heads, width) + plan = self._wrappers.get(key) + if plan is not None: + # Plans are usually created in capture warmup; pin them when the + # graph later records them so eager churn cannot free their buffers. + if not plan.pinned and torch.cuda.is_current_stream_capturing(): + plan.pinned = True + return plan + group = q_heads // kv_heads + if q_heads % kv_heads or group not in FLASHINFER_FA2_GROUP_SIZES: + raise ValueError( + f"flashinfer_fa2 tensor-core decode needs a query/KV head ratio in " + f"{FLASHINFER_FA2_GROUP_SIZES}, got {q_heads}/{kv_heads}" + ) + device = self.device + indptr = torch.zeros(rows + 1, dtype=torch.int32, device=device) + indices = torch.zeros(rows * width, dtype=torch.int32, device=device) + last_page_len = torch.full((rows,), self.page_size, dtype=torch.int32, device=device) + + def host(values): + tensor = torch.as_tensor(values, dtype=torch.int32) + return tensor.pin_memory() if torch.cuda.is_available() else tensor + + # Plan inputs live on the host: plan()/workspace_size() copy indptr and + # last_page_len to CPU, and a device source would synchronize each time. + capacity_indptr = host([row * width for row in range(rows + 1)]) + capacity_last = host([self.page_size] * rows) + capacity_indices = torch.zeros(rows * width, dtype=torch.int32) + if torch.cuda.is_available(): + capacity_indices = capacity_indices.pin_memory() + probe = torch.empty(256, dtype=torch.uint8, device=device) + wrapper = self._wrapper_cls( + probe, indptr, indices, last_page_len, "NHD", use_tensor_cores=True + ) + float_bytes, int_bytes = wrapper.workspace_size( + capacity_indptr, + capacity_indices, + capacity_last, + q_heads, + kv_heads, + self.head_dim, + self.page_size, + pos_encoding_mode="NONE", + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + disable_split_kv=True, + ) + wrapper.reset_workspace_buffer( + torch.empty(max(float_bytes, 256), dtype=torch.uint8, device=device), + torch.empty((int_bytes + 3) // 4, dtype=torch.int32, device=device), + ) + wrapper.plan( + capacity_indptr, + capacity_indices, + capacity_last, + q_heads, + kv_heads, + self.head_dim, + self.page_size, + pos_encoding_mode="NONE", + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + disable_split_kv=True, + ) + resolved = wrapper.resolved_backend + if resolved != "fa2": + raise RuntimeError( + f"flashinfer-python {self.info['version']} resolved the tensor-core " + f"wrapper to {resolved!r}; flashinfer_fa2 requires the FA2 kernel" + ) + self.info["resolved_backend"] = resolved + if len(self._wrappers) >= FLASHINFER_FA2_PLAN_CAPACITY: + # Captured plans are referenced by graph replays and must outlive + # them; only eager plans are evictable. + stale = next((name for name, entry in self._wrappers.items() if not entry.pinned), None) + if stale is not None: + del self._wrappers[stale] + plan = _Plan( + wrapper, + indptr, + indices, + last_page_len, + pinned=torch.cuda.is_current_stream_capturing(), + ) + self._plan_creations += 1 + self._wrappers[key] = plan + return plan + + def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): + if q.shape[0] == 0: + return torch.empty_like(q) + self._validate_metadata(block_tables, context_lengths) + rows, width = block_tables.shape + if rows != q.shape[0]: + raise ValueError( + f"flashinfer_fa2 needs one metadata row per query row, got {rows} rows " + f"for {q.shape[0]} queries" + ) + csr = self._primed.get((id(block_tables), id(context_lengths))) + if csr is None or not csr.owns(block_tables, context_lengths): + self._plan_misses += 1 + csr = self._csr_entry(block_tables, context_lengths) + self._write_csr(csr, block_tables, context_lengths) + plan = self._wrapper(rows, q.shape[1], key_cache.shape[2], width) + # One refresh per traversal (and one recorded group per capture), not + # per layer: every layer of a traversal shares the same plan buffers. + # Captures record their own copy so replay re-reads the staging buffers. + capturing = torch.cuda.is_current_stream_capturing() + synced = csr.captured if capturing else csr.synced + if plan not in synced: + plan.indptr.copy_(csr.indptr) + plan.indices[: csr.indices.numel()].copy_(csr.indices) + plan.last_page_len.copy_(csr.last_page_len) + synced.add(plan) + out = torch.empty_like(q) + plan.wrapper.run(q, (key_cache, value_cache), out=out) + out.mul_((context_lengths > 0).to(out.dtype).unsqueeze(-1).unsqueeze(-1)) + return out diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index 10149516..46137ade 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -286,6 +286,9 @@ def prepare_inputs(self, call, hidden): if self.last_event is not None: stream.wait_event(self.last_event) _copy_graph_inputs(entry, hidden, batch) + # Prime stateful kernels from the graph's private metadata, after the + # caller's tables land there and before capture/replay consumes them. + cache.prepare_attention(entry.metadata) if entry.graph is None: proxy = _DeviceCache(cache) entry.graph, entry.output = _capture( @@ -297,6 +300,8 @@ def prepare_inputs(self, call, hidden): ), ) self.captures += 1 + else: + self.cache.prepare_attention(batch) call.inputs_ready = True @torch.inference_mode() @@ -358,6 +363,7 @@ def prepare(self): bucketed=True, ) _copy_graph_inputs(entry, hidden, batch, count=0) + self.cache.prepare_attention(entry.metadata) entry.graph, entry.output = _capture( entry.hidden, self.capture_stream, @@ -463,6 +469,7 @@ def prepare_inputs(self, call, hidden, *, startup=False): stream.wait_event(self.last_event) if self.cache is not None: _copy_graph_inputs(entry, hidden, batch) + self.cache.prepare_attention(entry.metadata) else: entry.hidden[: len(hidden)].copy_(hidden) entry.hidden[len(hidden) :].zero_() @@ -480,6 +487,8 @@ def prepare_inputs(self, call, hidden, *, startup=False): entry.hidden, self.stream, self.pool, run, startup=startup ) self.captures += 1 + elif self.cache is not None: + self.cache.prepare_attention(batch) call.inputs_ready = True @torch.inference_mode() diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py index bad146d4..917435a9 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -207,6 +207,7 @@ def forward(self, batch: ForwardBatch) -> ForwardResult: batch.graphs.prepare_inputs(call, hidden) hidden, gate = batch.graphs.execute(call, hidden) else: + self.cache.prepare_attention(kv) hidden, gate = self.model.recurrent_prepared( hidden, kv, self.cache, compute_gate=batch.compute_gate )