From 69b7b1ceb8b3ddf076601af56d3d65cfa8dd1380 Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Sun, 1 Nov 2026 23:11:17 +0800 Subject: [PATCH 01/20] kernels: add FlashInfer FA2 paged decode for Ada SM89 needs FA2 because XQA rejects it and FlashAttention-2 rejects page size 16. Prime CSR metadata once per traversal and plan once per batch/table width; keep automatic selection unchanged and the backend explicit-only. --- pyproject.toml | 1 + vllm_rlt/attention/backend.py | 14 +- vllm_rlt/core/kv_cache_manager.py | 5 + vllm_rlt/entrypoints/cli.py | 1 + vllm_rlt/entrypoints/serve.py | 1 + vllm_rlt/kernels/flashinfer_fa2_attention.py | 177 +++++++++++++++++++ vllm_rlt/worker/executor.py | 1 + 7 files changed, 199 insertions(+), 1 deletion(-) create mode 100644 vllm_rlt/kernels/flashinfer_fa2_attention.py 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/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..f4cdc037 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() _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,11 @@ 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) + return AttentionBackend(kernel, _FLASHINFER_FA2_CAPABILITIES, dict(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..ea14980c --- /dev/null +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -0,0 +1,177 @@ +"""FlashInfer FA2 paged decode for Ada (SM89); no packed prefill or graphs. + +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__``. The plan is created once per (rows, heads, table width) with a +synthetic full-capacity table, so later steps only refresh device buffers. + +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. +""" + +from importlib.metadata import PackageNotFoundError, version + +import torch + +FLASHINFER_FA2_PAGE_SIZE = 16 +FLASHINFER_FA2_HEAD_DIM = 128 +FLASHINFER_FA2_WORKSPACE_BYTES = 128 * 1024 * 1024 + + +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 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 BatchDecodeWithPagedKVCacheWrapper + 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 = BatchDecodeWithPagedKVCacheWrapper + self.device = torch.device(device) + self.head_dim = head_dim + self.page_size = block_size + # Allocated on first execution so backend selection stays allocation-free. + self._workspace = None + self._wrappers = {} + self._primed = {} + self._plan_misses = 0 + self.info = dict( + backend="flashinfer_fa2", + package="flashinfer-python", + version=package_version, + kernel="fa2", + page_size=block_size, + head_dim=head_dim, + sm="8.9", + workspace_bytes=FLASHINFER_FA2_WORKSPACE_BYTES, + ) + + @property + def plan_misses(self): + return self._plan_misses + + @property + def plan_entries(self): + return len(self._wrappers) + + def prepare_step(self, batches): + """Build the CSR metadata of this traversal once, on the current stream.""" + 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 + 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") + if block_tables.shape[0] == 0: + continue + self._primed[id(block_tables)] = self._build_csr(block_tables, context_lengths) + + def _build_csr(self, block_tables, context_lengths): + rows, width = block_tables.shape + pages = (context_lengths.clamp_min(0) + self.page_size - 1) // self.page_size + indptr = torch.zeros(rows + 1, dtype=torch.int32, device=block_tables.device) + indptr[1:] = torch.cumsum(pages, dim=0) + columns = torch.arange(width, dtype=torch.int32, device=block_tables.device).unsqueeze(0) + indices = block_tables[columns < pages.unsqueeze(1)] + last_page_len = torch.where( + pages == 0, + torch.ones_like(pages), + context_lengths - (pages - 1) * self.page_size, + ).to(torch.int32) + return indptr, indices, last_page_len + + def _wrapper(self, rows, q_heads, kv_heads, width): + key = (rows, q_heads, kv_heads, width) + wrapper = self._wrappers.get(key) + if wrapper is not None: + return wrapper + if self._workspace is None: + self._workspace = torch.zeros( + FLASHINFER_FA2_WORKSPACE_BYTES, dtype=torch.uint8, device=self.device + ) + # Plan for the widest table this cache can present, so later steps only + # refresh metadata. Actual shorter tables are bounded by last_page_len. + indptr = torch.arange(rows + 1, dtype=torch.int32, device=self.device) * width + indices = torch.zeros(rows * width, dtype=torch.int32, device=self.device) + last_page_len = torch.full((rows,), self.page_size, dtype=torch.int32, device=self.device) + wrapper = self._wrapper_cls(self._workspace, "NHD", backend="fa2", use_tensor_cores=False) + wrapper.plan( + indptr, + indices, + last_page_len, + q_heads, + kv_heads, + self.head_dim, + self.page_size, + pos_encoding_mode="NONE", + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + for name in ( + "_paged_kv_indptr_buf", + "_paged_kv_indices_buf", + "_paged_kv_last_page_len_buf", + ): + if not hasattr(wrapper, name): + raise RuntimeError( + f"flashinfer-python {self.info['version']} no longer exposes {name}; " + "pin a supported version or update the adapter" + ) + self._wrappers[key] = wrapper + return wrapper + + def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): + if q.shape[0] == 0: + return torch.empty_like(q) + primed = self._primed.get(id(block_tables)) + if primed is None: + self._plan_misses += 1 + primed = self._build_csr(block_tables, context_lengths) + indptr, indices, last_page_len = primed + rows, width = block_tables.shape + wrapper = self._wrapper(rows, q.shape[1], key_cache.shape[2], width) + wrapper._paged_kv_indptr_buf.copy_(indptr) + wrapper._paged_kv_indices_buf[: indices.numel()].copy_(indices) + wrapper._paged_kv_last_page_len_buf.copy_(last_page_len) + out = torch.empty_like(q) + 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/executor.py b/vllm_rlt/worker/executor.py index bad146d4..a4d45bee 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -201,6 +201,7 @@ def forward(self, batch: ForwardBatch) -> ForwardResult: if hidden is None: raise ValueError("RECURRENT requires hidden state") for index, kv in enumerate(batch.metadata): + self.cache.prepare_attention(kv) if batch.graphs is not None: call = batch.graph_calls[index] if not call.inputs_ready: From 8b512b1dc20b36e302894110cc18e990c26c944f Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Sun, 1 Nov 2026 23:11:18 +0800 Subject: [PATCH 02/20] test: cover FlashInfer FA2 paged decode on SM89 Strided/GQA KV, padding tri-state, zero-length rows, batch composition, empty batches and tiny-Ouro token/exit-depth parity against Triton for LAST_EXITED and SHARED. --- tests/test_flashinfer_fa2_attention.py | 192 +++++++++++++++++++++++++ 1 file changed, 192 insertions(+) create mode 100644 tests/test_flashinfer_fa2_attention.py diff --git a/tests/test_flashinfer_fa2_attention.py b/tests/test_flashinfer_fa2_attention.py new file mode 100644 index 00000000..becc96c4 --- /dev/null +++ b/tests/test_flashinfer_fa2_attention.py @@ -0,0 +1,192 @@ +"""FlashInfer FA2 adapter: SM89 selection, CSR metadata and eager decode.""" + +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 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 + 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 not any( + ( + capabilities.packed_prefill, + capabilities.cuda_graphs, + capabilities.async_scheduling, + capabilities.device_context_lengths, + ) + ) + + +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) + + +@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 + + attention.prepare_step((SimpleNamespace(block_tables=tables[:3], context_lengths=lengths[:3]),)) + smaller = attention(q[:3], keys, values, tables[:3], lengths[:3]) + 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 +@pytest.mark.parametrize("layout", ["last_exited", "shared"]) +def test_tiny_ouro_matches_triton(layout): + torch.manual_seed(123) + model = OuroForCausalLM(_tiny_config()).to(device="cuda", dtype=torch.bfloat16) + traces = {str(i): [4, 2 + i % 3, 4, 2, 3, 4] for i in range(3)} + outputs = [] + for backend in ("triton", "flashinfer_fa2"): + 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=ExecutionConfig(), + exit_config=ExitConfig("trace", depths_by_request=traces), + attention_backend=backend, + ) + 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 + outputs.append(finished) + assert outputs[0] == outputs[1] + kernel = engine.cache_manager.attention.kernel + assert kernel.plan_misses == 0 From 31d4a36886ab671e9e2f73e9d5a3bad0d68271eb Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Sun, 1 Nov 2026 23:11:20 +0800 Subject: [PATCH 03/20] bench: select the native attention backend for GSM8K --- benchmarks/gsm8k.py | 17 ++++++++++++++++- benchmarks/gsm8k_backends.py | 12 ++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index d1ecc7f0..c2abe001 100644 --- a/benchmarks/gsm8k.py +++ b/benchmarks/gsm8k.py @@ -589,8 +589,10 @@ def run(args): raise RuntimeError("Run with one scheduler-assigned GPU") output = Path(args.output) output.mkdir(parents=True, exist_ok=False) + attention_backend = getattr(args, "attention_backend", "triton") metadata = { "backend": args.backend, + "attention_backend": attention_backend, "protocol_sha256": file_digest(args.protocol), "source": source(), "cuda_visible_devices": os.environ["CUDA_VISIBLE_DEVICES"], @@ -609,7 +611,14 @@ 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, + ) load_seconds = time.monotonic() - start rows = [] try: @@ -785,6 +794,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") diff --git a/benchmarks/gsm8k_backends.py b/benchmarks/gsm8k_backends.py index 15d57165..283db2fc 100644 --- a/benchmarks/gsm8k_backends.py +++ b/benchmarks/gsm8k_backends.py @@ -101,7 +101,15 @@ 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", + ): 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,7 +146,7 @@ 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), From 66e0734de0cd2a7a7a902879432378647bcb9037 Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Sun, 1 Nov 2026 23:11:21 +0800 Subject: [PATCH 04/20] docs: document the FlashInfer FA2 Ada backend --- README.md | 4 +- docs/README.md | 1 + docs/accuracy.md | 18 +++++++ docs/flashinfer-fa2.md | 105 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 docs/flashinfer-fa2.md diff --git a/README.md b/README.md index 00ed4e17..3c87d30a 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; 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/docs/README.md b/docs/README.md index b0ab7777..8f996c5c 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 paged decode, explicit opt-in, validation scope and limitations | | [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..0d672641 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -198,6 +198,24 @@ 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`. The flag is recorded in the +summary and does not change the frozen baseline fingerprint. + +On one L20 (SM89, torch 2.8.0, FlashInfer 0.7.2.post1), a custom 10-question +subset produced **3/10 with Triton and 3/10 with FlashInfer FA2**; 8/10 generated +token sequences were identical, and the two divergences were on questions both +backends answered incorrectly. The first divergences were near ties (top1/top2 +margins 0.0–0.0625), consistent with BF16 kernel accumulation order rather than an +adapter defect. This is a diagnostic subset, not the GSM8K-87 gate. + +The stored GSM8K-87 baseline fingerprints `torch==2.13.0`; a run on another torch +version cannot reuse its 59/87 baseline. Run the default protocol on a torch-2.13 +environment before quoting an 87-question score for either backend. + + 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..b24cdba4 --- /dev/null +++ b/docs/flashinfer-fa2.md @@ -0,0 +1,105 @@ +# FlashInfer FA2 backend (Ada / SM89) + +`--attention-backend flashinfer_fa2` runs Ouro's paged decode through +FlashInfer's **FA2** `BatchDecodeWithPagedKVCacheWrapper`. 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. +- Synchronous eager execution. +- No CUDA Graphs, async scheduling, packed prefill, PD 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(), +) as llm: + ... +``` + +or `--attention-backend flashinfer_fa2` on the offline CLI/server. + +## Design + +- `prepare_attention` / `ModelRuntime.forward` prime the backend once per + traversal (`KVCacheManager.prepare_attention`), so the CSR metadata + (`indptr`, `indices`, `last_page_len`) is built once per depth view and reused + across all decoder layers. +- The FA2 plan is created once per `(rows, q_heads, kv_heads, table_width)` with + a synthetic full-capacity table; later steps only refresh device buffers. +- 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. +- Padding rows (`context_lengths == 0`) are masked to exact zeros. +- No `prefill` entry point: the default Ouro path never calls it, and + `packed_prefill` stays false. + +## Validation (L20, 2026-11-01) + +Environment: 1×L20 (sm_89, 46 GB, driver 580.126.20), torch 2.8.0, FlashInfer +0.7.2.post1, Triton 3.7.1. The L20 pod's CUDA 12.1 needed a local removal of +FlashInfer's CUDA-12.9-only `--compress-mode=size` flag; use CUDA 12.9 for a +clean build. + +- Unit tests: `tests/test_flashinfer_fa2_attention.py` covers selection, + rejection paths, strided/GQA KV, padding tri-state, zero-length rows, batch + composition, empty batches, and tiny-Ouro token/exit-depth parity against + Triton for LAST_EXITED and SHARED. +- Real Ouro-1.4B: a 10-question custom GSM8K subset produced the same correct + count (3/10) with Triton and FA2; 8/10 generated token sequences were + identical. The first divergences sit on questions both backends answered + incorrectly, at near-tied logits (top1/top2 margins 0.0–0.0625). +- The official GSM8K-87 fingerprint requires torch 2.13.0; the L20 pod ran + torch 2.8, so the stored 59/87 baseline was not claimed. Run the official + protocol on a torch-2.13 environment before quoting a 87-question score. + +## Performance (diagnostic, not a benchmark) + +Kernel-only, Ouro attention shapes, page 16, random KV, single L20: + +| B | ctx | FA2 ms | Triton ms | FA2/Triton | +|---:|---:|---:|---:|---:| +| 1 | 256 | 0.0225 | 0.0584 | 0.39 | +| 1 | 2048 | 0.0334 | 0.0609 | 0.55 | +| 1 | 8192 | 0.0243 | 0.0587 | 0.41 | +| 8 | 2048 | 0.1856 | 0.1933 | 0.96 | +| 32 | 8192 | 2.9407 | 2.9881 | 0.98 | + +FA2 is ~2–3× faster at low concurrency/short context and at parity at high +concurrency/long context. End-to-end, the 10-question real-model run took +84.8 s with Triton and 79.7 s with FA2; the model forward dominates, so the +attention gain is small. Do not treat these as production speedup claims. + +## Follow-ups + +- CUDA Graph: prime metadata before capture and enable `cuda_graphs`. +- Async scheduling/device-owned metadata. +- Auto-selection on SM89 after broader validation. +- More page sizes/head dimensions, packed prefill, PD and speculative paths. +- Reconcile with #112 if XQA lands first; this backend keeps a distinct + `flashinfer_fa2` name and module. From 31c9f82406abcc5dbd48912834b733a30fb71148 Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Mon, 2 Nov 2026 00:06:18 +0800 Subject: [PATCH 05/20] docs: record L20 GSM8K-87 results for FlashInfer FA2 Official protocol: Triton 61/87 and FA2 61/87, both passing the 59/87 baseline; 65/87 token-identical and two cancelling correct-answer flips. FA2 has no demonstrated end-to-end speedup on this workload. --- docs/accuracy.md | 24 ++++++++++++++++-------- docs/flashinfer-fa2.md | 36 +++++++++++++++++++++--------------- 2 files changed, 37 insertions(+), 23 deletions(-) diff --git a/docs/accuracy.md b/docs/accuracy.md index 0d672641..b253c78a 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -204,16 +204,24 @@ above are explicit project choices rather than an exact paper reproduction. for the native backend; the default stays `triton`. The flag is recorded in the summary and does not change the frozen baseline fingerprint. -On one L20 (SM89, torch 2.8.0, FlashInfer 0.7.2.post1), a custom 10-question -subset produced **3/10 with Triton and 3/10 with FlashInfer FA2**; 8/10 generated -token sequences were identical, and the two divergences were on questions both -backends answered incorrectly. The first divergences were near ties (top1/top2 -margins 0.0–0.0625), consistent with BF16 kernel accumulation order rather than an -adapter defect. This is a diagnostic subset, not the GSM8K-87 gate. +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: + +| backend | correct | accuracy | gate | +|---|---:|---:|---| +| triton | 61/87 | 70.11% | pass (+2.30 pp) | +| flashinfer_fa2 | 61/87 | 70.11% | pass (+2.30 pp) | + +65/87 generated token sequences are identical. Two correct-answer flips cancel: +question 360 goes incorrect → correct with FA2, and question 426 goes +correct → incorrect. FA2 has one length-limited runaway (question 182 reaches +1024 tokens; Triton stops at 40), so its mean generated length is 115.9 vs +103.7 tokens and its generation time is 732.6 s vs 637.4 s. This backend does +not demonstrate an end-to-end speedup on this workload. The stored GSM8K-87 baseline fingerprints `torch==2.13.0`; a run on another torch -version cannot reuse its 59/87 baseline. Run the default protocol on a torch-2.13 -environment before quoting an 87-question score for either backend. +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 diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index b24cdba4..0288201b 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -59,24 +59,27 @@ or `--attention-backend flashinfer_fa2` on the offline CLI/server. - No `prefill` entry point: the default Ouro path never calls it, and `packed_prefill` stays false. -## Validation (L20, 2026-11-01) +## Validation (L20, 2026-11-01/02) -Environment: 1×L20 (sm_89, 46 GB, driver 580.126.20), torch 2.8.0, FlashInfer -0.7.2.post1, Triton 3.7.1. The L20 pod's CUDA 12.1 needed a local removal of -FlashInfer's CUDA-12.9-only `--compress-mode=size` flag; use CUDA 12.9 for a -clean build. +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. - Unit tests: `tests/test_flashinfer_fa2_attention.py` covers selection, rejection paths, strided/GQA KV, padding tri-state, zero-length rows, batch composition, empty batches, and tiny-Ouro token/exit-depth parity against Triton for LAST_EXITED and SHARED. -- Real Ouro-1.4B: a 10-question custom GSM8K subset produced the same correct - count (3/10) with Triton and FA2; 8/10 generated token sequences were - identical. The first divergences sit on questions both backends answered - incorrectly, at near-tied logits (top1/top2 margins 0.0–0.0625). -- The official GSM8K-87 fingerprint requires torch 2.13.0; the L20 pod ran - torch 2.8, so the stored 59/87 baseline was not claimed. Run the official - protocol on a torch-2.13 environment before quoting a 87-question score. +- Official GSM8K-87 (the default protocol with its 59/87 baseline): + **Triton 61/87, FlashInfer FA2 61/87**; both pass the +2.30 pp gate. + - 65/87 generated token sequences are identical. + - Two correct-answer flips cancel: question 360 changes from incorrect + (Triton) to correct (FA2); question 426 changes from correct to incorrect. + - FA2 has one length-limited runaway (question 182 reaches 1024 tokens; Triton + stops at 40), which inflates its generation time and mean length. +- 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) @@ -91,9 +94,12 @@ Kernel-only, Ouro attention shapes, page 16, random KV, single L20: | 32 | 8192 | 2.9407 | 2.9881 | 0.98 | FA2 is ~2–3× faster at low concurrency/short context and at parity at high -concurrency/long context. End-to-end, the 10-question real-model run took -84.8 s with Triton and 79.7 s with FA2; the model forward dominates, so the -attention gain is small. Do not treat these as production speedup claims. +concurrency/long context. End-to-end on the official GSM8K-87 run, FA2 took +732.6 s and Triton 637.4 s, with mean generated lengths 115.9 and 103.7 tokens; +the 1024-token runaway accounts for much of the gap. The model forward dominates, +so this backend has **no demonstrated end-to-end speedup** in this workload. Do +not treat these as production speedup claims. + ## Follow-ups From 03b3089076d68f88cf24e41ed1898cb056e6e869 Mon Sep 17 00:00:00 2001 From: sysuTobo Date: Mon, 2 Nov 2026 00:23:06 +0800 Subject: [PATCH 06/20] docs: quantify L20 FA2 tie flips and matched throughput Record the first-divergence margins for the two GSM8K correct flips and q182, and the matched C1/C8/C32 fixed-work timings (FA2/Triton 1.05/0.97/1.03, within noise). --- docs/accuracy.md | 11 +++++++---- docs/flashinfer-fa2.md | 21 ++++++++++++++++++--- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/docs/accuracy.md b/docs/accuracy.md index b253c78a..9fc1e348 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -214,10 +214,13 @@ On one L20 (SM89, torch 2.13.0, CUDA 13.0, FlashInfer 0.7.2.post1, Triton 65/87 generated token sequences are identical. Two correct-answer flips cancel: question 360 goes incorrect → correct with FA2, and question 426 goes -correct → incorrect. FA2 has one length-limited runaway (question 182 reaches -1024 tokens; Triton stops at 40), so its mean generated length is 115.9 vs -103.7 tokens and its generation time is 732.6 s vs 637.4 s. This backend does -not demonstrate an end-to-end speedup on this workload. +correct → incorrect. The first divergences are near ties (logits maxabs +<= 0.17, top1/top2 margins 0–0.125), consistent with BF16 accumulation order. +FA2 has one length-limited runaway (question 182 reaches 1024 tokens; Triton +stops at 40), so its mean generated length is 115.9 vs 103.7 tokens and its +generation time is 732.6 s vs 637.4 s. A matched 64-token fixed-work run at +concurrency 1/8/32 gives FA2/Triton time ratios 1.05/0.97/1.03, within noise; +this backend does not demonstrate an end-to-end speedup on this workload. 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 diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index 0288201b..2fa31604 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -76,6 +76,10 @@ official run used the pip `cuda-toolkit[nvcc]==13.0.3` tree. - 65/87 generated token sequences are identical. - Two correct-answer flips cancel: question 360 changes from incorrect (Triton) to correct (FA2); question 426 changes from correct to incorrect. + - The flips are BF16 ties: first-divergence logits differ by <= 0.17. q360 + has margins 0.0 (Triton) vs 0.125 (FA2); q426 is the reverse; q182's first + token is tied at 0.0 (FA2) vs 0.125 (Triton), which sends FA2 into the + 1024-token runaway. - FA2 has one length-limited runaway (question 182 reaches 1024 tokens; Triton stops at 40), which inflates its generation time and mean length. - An earlier 10-question custom subset on torch 2.8 was used as a smoke: @@ -96,9 +100,20 @@ Kernel-only, Ouro attention shapes, page 16, random KV, single L20: FA2 is ~2–3× faster at low concurrency/short context and at parity at high concurrency/long context. End-to-end on the official GSM8K-87 run, FA2 took 732.6 s and Triton 637.4 s, with mean generated lengths 115.9 and 103.7 tokens; -the 1024-token runaway accounts for much of the gap. The model forward dominates, -so this backend has **no demonstrated end-to-end speedup** in this workload. Do -not treat these as production speedup claims. +the 1024-token runaway accounts for much of the gap. + +A matched fixed-work run (64 output tokens, same prompts, two repeats per cell) +gives: + +| concurrency | Triton s | FA2 s | FA2/Triton | +|---:|---:|---:|---:| +| 1 | 4.37 | 4.59 | 1.05 | +| 8 | 6.62 | 6.41 | 0.97 | +| 32 | 11.44 | 11.79 | 1.03 | + +The differences are within run-to-run noise. The model forward dominates, so +this backend has **no demonstrated end-to-end speedup** in this workload. Do not +treat these as production speedup claims. ## Follow-ups From 8a71defa074cd0d98d7aa8525025827a2432ce2c Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 02:07:18 +0800 Subject: [PATCH 07/20] feat: enable CUDA Graphs and async steps in FlashInfer FA2 --- tests/test_flashinfer_fa2_attention.py | 191 ++++++++++--- vllm_rlt/attention/backend.py | 2 +- vllm_rlt/kernels/flashinfer_fa2_attention.py | 283 ++++++++++++++----- vllm_rlt/worker/cuda_graph.py | 9 + vllm_rlt/worker/executor.py | 2 +- 5 files changed, 375 insertions(+), 112 deletions(-) diff --git a/tests/test_flashinfer_fa2_attention.py b/tests/test_flashinfer_fa2_attention.py index becc96c4..c4e93342 100644 --- a/tests/test_flashinfer_fa2_attention.py +++ b/tests/test_flashinfer_fa2_attention.py @@ -14,7 +14,7 @@ SamplingParams, SchedulerConfig, ) -from vllm_rlt.attention import create_backend, validate_backend_name +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 @@ -33,6 +33,7 @@ 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) @@ -46,14 +47,28 @@ def test_backend_name_and_explicit_selection(fake_flashinfer): selected = create_backend("flashinfer_fa2", torch.device("cuda"), torch.bfloat16, 128, 16) assert selected.info["selected_backend"] == "flashinfer_fa2" capabilities = selected.capabilities - assert not any( - ( - capabilities.packed_prefill, - capabilities.cuda_graphs, - capabilities.async_scheduling, - capabilities.device_context_lengths, - ) + 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): @@ -109,8 +124,11 @@ def test_strided_gqa_padding_and_zero_rows(): assert torch.equal(out[:3], actual[:3]) assert torch.count_nonzero(out[3]) == 0 - attention.prepare_step((SimpleNamespace(block_tables=tables[:3], context_lengths=lengths[:3]),)) - smaller = attention(q[:3], keys, values, tables[:3], lengths[:3]) + 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 @@ -153,40 +171,127 @@ def _tiny_config(): ) +@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, 145], 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] + assert csr.last_page_len.tolist() == [ + int(length) - (count - 1) * 16 for length, count in zip(lengths.tolist(), pages) + ] + + +@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 + + +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"]) -def test_tiny_ouro_matches_triton(layout): +@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) - traces = {str(i): [4, 2 + i % 3, 4, 2, 3, 4] for i in range(3)} - outputs = [] - for backend in ("triton", "flashinfer_fa2"): - 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=ExecutionConfig(), - exit_config=ExitConfig("trace", depths_by_request=traces), - attention_backend=backend, - ) - 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 - outputs.append(finished) - assert outputs[0] == outputs[1] - kernel = engine.cache_manager.attention.kernel - assert kernel.plan_misses == 0 + 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 diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index f4cdc037..1c0fb021 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -66,7 +66,7 @@ def generation(self): device_context_lengths=True, ), } -_FLASHINFER_FA2_CAPABILITIES = BackendCapabilities() +_FLASHINFER_FA2_CAPABILITIES = BackendCapabilities(cuda_graphs=True, async_scheduling=True) _AUTO_PRIORITY = ("flash_attn_4", "flash_attn_3", "flash_attn_2", "triton", "torch") diff --git a/vllm_rlt/kernels/flashinfer_fa2_attention.py b/vllm_rlt/kernels/flashinfer_fa2_attention.py index ea14980c..110b423b 100644 --- a/vllm_rlt/kernels/flashinfer_fa2_attention.py +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -1,21 +1,46 @@ -"""FlashInfer FA2 paged decode for Ada (SM89); no packed prefill or graphs. +"""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__``. The plan is created once per (rows, heads, table width) with a -synthetic full-capacity table, so later steps only refresh device buffers. +``__call__``. Plans are created once per (rows, heads, table width) at the +cache's full table capacity, so later steps only refresh device buffers. + +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_WORKSPACE_BYTES = 128 * 1024 * 1024 +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): @@ -27,6 +52,46 @@ def _version_tuple(value): 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") + + 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) + + 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 ``__call__`` copies it in, so replays see new metadata. + """ + + __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: @@ -49,20 +114,20 @@ def __init__(self, device, dtype, head_dim, block_size): "flashinfer_fa2 requires flashinfer-python>=0.7.2,<0.8; " f"found {package_version}" ) - from flashinfer.decode import BatchDecodeWithPagedKVCacheWrapper + 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 = BatchDecodeWithPagedKVCacheWrapper + self._wrapper_cls = CUDAGraphBatchDecodeWithPagedKVCacheWrapper self.device = torch.device(device) self.head_dim = head_dim self.page_size = block_size - # Allocated on first execution so backend selection stays allocation-free. - self._workspace = None self._wrappers = {} + self._column_index = {} + self._csr = {} self._primed = {} self._plan_misses = 0 self.info = dict( @@ -70,10 +135,13 @@ def __init__(self, device, dtype, head_dim, block_size): package="flashinfer-python", version=package_version, kernel="fa2", + tensor_cores=True, + cuda_graph_wrapper=True, page_size=block_size, head_dim=head_dim, sm="8.9", - workspace_bytes=FLASHINFER_FA2_WORKSPACE_BYTES, + cuda_graphs=True, + async_scheduling=True, ) @property @@ -85,58 +153,128 @@ def plan_entries(self): return len(self._wrappers) def prepare_step(self, batches): - """Build the CSR metadata of this traversal once, on the current stream.""" + """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 - 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") + self._validate_metadata(block_tables, context_lengths) if block_tables.shape[0] == 0: continue - self._primed[id(block_tables)] = self._build_csr(block_tables, context_lengths) + csr = self._csr_entry(block_tables, context_lengths) + self._write_csr(csr, block_tables, context_lengths) + self._primed[(id(block_tables), id(context_lengths))] = csr - def _build_csr(self, block_tables, context_lengths): - rows, width = block_tables.shape - pages = (context_lengths.clamp_min(0) + self.page_size - 1) // self.page_size - indptr = torch.zeros(rows + 1, dtype=torch.int32, device=block_tables.device) - indptr[1:] = torch.cumsum(pages, dim=0) - columns = torch.arange(width, dtype=torch.int32, device=block_tables.device).unsqueeze(0) - indices = block_tables[columns < pages.unsqueeze(1)] - last_page_len = torch.where( + @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 + csr.indptr[0] = 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) * self.page_size, - ).to(torch.int32) - return indptr, indices, last_page_len + context_lengths - (pages - 1) * page, + out=csr.last_page_len, + ) def _wrapper(self, rows, q_heads, kv_heads, width): key = (rows, q_heads, kv_heads, width) - wrapper = self._wrappers.get(key) - if wrapper is not None: - return wrapper - if self._workspace is None: - self._workspace = torch.zeros( - FLASHINFER_FA2_WORKSPACE_BYTES, dtype=torch.uint8, device=self.device + 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}" ) - # Plan for the widest table this cache can present, so later steps only - # refresh metadata. Actual shorter tables are bounded by last_page_len. - indptr = torch.arange(rows + 1, dtype=torch.int32, device=self.device) * width - indices = torch.zeros(rows * width, dtype=torch.int32, device=self.device) - last_page_len = torch.full((rows,), self.page_size, dtype=torch.int32, device=self.device) - wrapper = self._wrapper_cls(self._workspace, "NHD", backend="fa2", use_tensor_cores=False) + 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) + # Plan for the widest table this cache can present; live CSR values are + # written before every run and unused splits are masked by the wrapper. + capacity_indptr = torch.arange(rows + 1, dtype=torch.int32, device=device) * width + capacity_indices = torch.zeros(rows * width, dtype=torch.int32, device=device) + capacity_last = torch.full((rows,), self.page_size, dtype=torch.int32, device=device) + 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, + ) + 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( - indptr, - indices, - last_page_len, + capacity_indptr, + capacity_indices, + capacity_last, q_heads, kv_heads, self.head_dim, @@ -145,33 +283,44 @@ def _wrapper(self, rows, q_heads, kv_heads, width): q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16, ) - for name in ( - "_paged_kv_indptr_buf", - "_paged_kv_indices_buf", - "_paged_kv_last_page_len_buf", - ): - if not hasattr(wrapper, name): - raise RuntimeError( - f"flashinfer-python {self.info['version']} no longer exposes {name}; " - "pin a supported version or update the adapter" - ) - self._wrappers[key] = wrapper - return wrapper + 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._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) - primed = self._primed.get(id(block_tables)) - if primed is None: - self._plan_misses += 1 - primed = self._build_csr(block_tables, context_lengths) - indptr, indices, last_page_len = primed + self._validate_metadata(block_tables, context_lengths) rows, width = block_tables.shape - wrapper = self._wrapper(rows, q.shape[1], key_cache.shape[2], width) - wrapper._paged_kv_indptr_buf.copy_(indptr) - wrapper._paged_kv_indices_buf[: indices.numel()].copy_(indices) - wrapper._paged_kv_last_page_len_buf.copy_(last_page_len) + 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) + # Copied per layer on purpose: under CUDA Graph capture these nodes are + # recorded and re-read the refreshed staging buffers on every replay. + plan.indptr.copy_(csr.indptr) + plan.indices.copy_(csr.indices) + plan.last_page_len.copy_(csr.last_page_len) out = torch.empty_like(q) - wrapper.run(q, (key_cache, value_cache), out=out) + 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 a4d45bee..917435a9 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -201,13 +201,13 @@ def forward(self, batch: ForwardBatch) -> ForwardResult: if hidden is None: raise ValueError("RECURRENT requires hidden state") for index, kv in enumerate(batch.metadata): - self.cache.prepare_attention(kv) if batch.graphs is not None: call = batch.graph_calls[index] if not call.inputs_ready: 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 ) From 23e28f84970397284e64b1f1f355008b7f40428c Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 02:07:19 +0800 Subject: [PATCH 08/20] bench: expose --cuda-graphs in the GSM8K harness --- benchmarks/gsm8k.py | 10 ++++++++++ benchmarks/gsm8k_backends.py | 5 ++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index c2abe001..b156b2f4 100644 --- a/benchmarks/gsm8k.py +++ b/benchmarks/gsm8k.py @@ -590,9 +590,13 @@ def run(args): output = Path(args.output) output.mkdir(parents=True, exist_ok=False) attention_backend = getattr(args, "attention_backend", "triton") + cuda_graphs = getattr(args, "cuda_graphs", False) + if cuda_graphs and args.backend != "native": + raise ValueError("--cuda-graphs applies to the native backend only") 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"], @@ -618,6 +622,7 @@ def run(args): protocol["max_length"], exit_policy, attention_backend=attention_backend, + cuda_graphs=cuda_graphs, ) load_seconds = time.monotonic() - start rows = [] @@ -811,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 283db2fc..8cc73fda 100644 --- a/benchmarks/gsm8k_backends.py +++ b/benchmarks/gsm8k_backends.py @@ -109,6 +109,7 @@ def __init__( 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). @@ -151,7 +152,9 @@ def __init__( 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( From b7411f914bde31b3f8a9a44b7d897047038ae55f Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 02:07:21 +0800 Subject: [PATCH 09/20] docs: record L20 FA2 CUDA Graph and async validation --- README.md | 4 +- docs/README.md | 2 +- docs/accuracy.md | 35 ++++++----- docs/flashinfer-fa2.md | 140 ++++++++++++++++++++++++----------------- 4 files changed, 108 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 3c87d30a..4fc95ae5 100644 --- a/README.md +++ b/README.md @@ -92,8 +92,8 @@ in mixed prefill/decode batches. 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. FlashInfer FA2 paged decode - is an explicit opt-in backend for Ada (SM89) GPUs; see - [the FA2 guide](docs/flashinfer-fa2.md). + 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/docs/README.md b/docs/README.md index 8f996c5c..d830cb29 100644 --- a/docs/README.md +++ b/docs/README.md @@ -24,7 +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 paged decode, explicit opt-in, validation scope and limitations | +| [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 9fc1e348..9a2b6cf6 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -201,26 +201,33 @@ 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`. The flag is recorded in the -summary and does not change the frozen baseline fingerprint. +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: -| backend | correct | accuracy | gate | +| run | correct | accuracy | gate | |---|---:|---:|---| | triton | 61/87 | 70.11% | pass (+2.30 pp) | -| flashinfer_fa2 | 61/87 | 70.11% | pass (+2.30 pp) | - -65/87 generated token sequences are identical. Two correct-answer flips cancel: -question 360 goes incorrect → correct with FA2, and question 426 goes -correct → incorrect. The first divergences are near ties (logits maxabs -<= 0.17, top1/top2 margins 0–0.125), consistent with BF16 accumulation order. -FA2 has one length-limited runaway (question 182 reaches 1024 tokens; Triton -stops at 40), so its mean generated length is 115.9 vs 103.7 tokens and its -generation time is 732.6 s vs 637.4 s. A matched 64-token fixed-work run at -concurrency 1/8/32 gives FA2/Triton time ratios 1.05/0.97/1.03, within noise; -this backend does not demonstrate an end-to-end speedup on this workload. +| flashinfer_fa2 eager | 60/87 | 68.97% | pass (+1.15 pp) | +| flashinfer_fa2 + CUDA Graphs | 59/87 | 67.82% | pass (+0.00 pp) | + +Eager and graph FA2 share 69/87 generated token sequences. Padded graph rows +and the full-width plan change the tensor-core split schedule, so 18 first +divergences are BF16 ties (logits maxabs <= 0.17, top1/top2 margins 0–0.125) +and each one changes the whole continuation. The graphs run has no +length-limited runaway (mean 106.7 tokens); the eager run reaches 1024 tokens +on question 182 (mean 117.1), and Triton stops that question at 40. + +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: a matched 64-token fixed-work run went from 5.24/5.78/6.22 s +(eager) to 1.67/1.92/2.32 s (graphs) at concurrency 1/8/32. That is a launch- +overhead result on one L20, not a production throughput qualification; +async+graphs is within noise of graphs. 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 diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index 2fa31604..7aaef90e 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -1,17 +1,17 @@ # FlashInfer FA2 backend (Ada / SM89) `--attention-backend flashinfer_fa2` runs Ouro's paged decode through -FlashInfer's **FA2** `BatchDecodeWithPagedKVCacheWrapper`. 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. +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. -- Synchronous eager execution. -- No CUDA Graphs, async scheduling, packed prefill, PD or speculative paths. +- 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 @@ -37,27 +37,50 @@ with LLM( 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(), + execution_config=ExecutionConfig(cuda_graphs=True, async_scheduling=True), ) as llm: ... ``` -or `--attention-backend flashinfer_fa2` on the offline CLI/server. +or `--attention-backend flashinfer_fa2` on the offline CLI/server, combined with +the engine's ordinary CUDA Graph and asynchronous-scheduling options. ## Design -- `prepare_attention` / `ModelRuntime.forward` prime the backend once per - traversal (`KVCacheManager.prepare_attention`), so the CSR metadata - (`indptr`, `indices`, `last_page_len`) is built once per depth view and reused - across all decoder layers. -- The FA2 plan is created once per `(rows, q_heads, kv_heads, table_width)` with - a synthetic full-capacity table; later steps only refresh device buffers. +- **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 full table capacity but takes the schedule + from the live CSR buffers and masks unused splits, which is the graph + contract FlashInfer documents for this wrapper. +- **Plan cache.** One plan per `(rows, q_heads, kv_heads, table_width)`, created + on first use with a synthetic full-capacity table. Each plan owns its CSR + buffers and a workspace sized by `workspace_size(...)`; FlashInfer keeps + plan-time tables in that workspace, so plans cannot share it. Plans captured + by a CUDA Graph are pinned; eager plans are evicted first when the cache + (32 entries) is full. Every plan is created outside graph capture. +- **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, + so a step with device-owned context lengths does not read back to the host. +- **Replay refreshes the plan buffers.** `__call__` copies the staging CSR into + the plan's buffers; graph capture records those copies, 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. -- Padding rows (`context_lengths == 0`) are masked to exact zeros. - No `prefill` entry point: the default Ouro path never calls it, and - `packed_prefill` stays false. + `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) @@ -67,60 +90,65 @@ 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. -- Unit tests: `tests/test_flashinfer_fa2_attention.py` covers selection, - rejection paths, strided/GQA KV, padding tri-state, zero-length rows, batch - composition, empty batches, and tiny-Ouro token/exit-depth parity against - Triton for LAST_EXITED and SHARED. -- Official GSM8K-87 (the default protocol with its 59/87 baseline): - **Triton 61/87, FlashInfer FA2 61/87**; both pass the +2.30 pp gate. - - 65/87 generated token sequences are identical. - - Two correct-answer flips cancel: question 360 changes from incorrect - (Triton) to correct (FA2); question 426 changes from correct to incorrect. - - The flips are BF16 ties: first-divergence logits differ by <= 0.17. q360 - has margins 0.0 (Triton) vs 0.125 (FA2); q426 is the reverse; q182's first - token is tied at 0.0 (FA2) vs 0.125 (Triton), which sends FA2 into the - 1024-token runaway. - - FA2 has one length-limited runaway (question 182 reaches 1024 tokens; Triton - stops at 40), which inflates its generation time and mean length. +- `tests/test_flashinfer_fa2_attention.py` (15 GPU cases): 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; adapter-level CUDA Graph capture/replay + with refreshed metadata; and tiny-Ouro engine 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) | 60/87 | 68.97% | pass (+1.15 pp) | +| FA2 + CUDA Graphs | 59/87 | 67.82% | pass (+0.00 pp) | + + Eager and graphs share **69/87** token sequences; 18 diverge at a BF16 tie + because padded graph rows and the full-width plan change the tensor-core + split schedule, and each divergence changes the whole continuation. The + graphs run has no length-limited runaway (mean 106.7 tokens); the eager run + hits one (question 182, mean 117.1 tokens). +- 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) -Kernel-only, Ouro attention shapes, page 16, random KV, single L20: +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: -| B | ctx | FA2 ms | Triton ms | FA2/Triton | +| concurrency | Triton eager | FA2 eager | FA2 graphs | FA2 async+graphs | |---:|---:|---:|---:|---:| -| 1 | 256 | 0.0225 | 0.0584 | 0.39 | -| 1 | 2048 | 0.0334 | 0.0609 | 0.55 | -| 1 | 8192 | 0.0243 | 0.0587 | 0.41 | -| 8 | 2048 | 0.1856 | 0.1933 | 0.96 | -| 32 | 8192 | 2.9407 | 2.9881 | 0.98 | +| 1 | 4.94 s (13.0 tok/s) | 5.24 s (12.2) | 1.67 s (38.4) | 1.66 s (38.7) | +| 8 | 5.50 s (93.1) | 5.78 s (88.6) | 1.92 s (267.4) | 1.88 s (272.1) | +| 32 | 6.12 s (334.5) | 6.22 s (329.4) | 2.32 s (881.7) | 2.23 s (918.8) | -FA2 is ~2–3× faster at low concurrency/short context and at parity at high -concurrency/long context. End-to-end on the official GSM8K-87 run, FA2 took -732.6 s and Triton 637.4 s, with mean generated lengths 115.9 and 103.7 tokens; -the 1024-token runaway accounts for much of the gap. - -A matched fixed-work run (64 output tokens, same prompts, two repeats per cell) -gives: - -| concurrency | Triton s | FA2 s | FA2/Triton | -|---:|---:|---:|---:| -| 1 | 4.37 | 4.59 | 1.05 | -| 8 | 6.62 | 6.41 | 0.97 | -| 32 | 11.44 | 11.79 | 1.03 | - -The differences are within run-to-run noise. The model forward dominates, so -this backend has **no demonstrated end-to-end speedup** in this workload. Do not -treat these as production speedup claims. +Eager FA2 is at parity with Triton; CUDA Graphs remove most per-step launch +overhead on this small model (roughly 2.6–3.1× over FA2 eager and 3.0–3.4× +over Triton eager). The official GSM8K-87 runs agree in direction (eager +784.6 s, graphs 281.5 s, with different generated lengths). Asynchronous +scheduling adds no measurable gain at these batch sizes. These are single-GPU, +single-wave diagnostics, not production throughput claims. +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 -- CUDA Graph: prime metadata before capture and enable `cuda_graphs`. -- Async scheduling/device-owned metadata. - 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. +- Size plan workspaces per graph-entry width; full-capacity plans allocate the + widest table even for short contexts. - Reconcile with #112 if XQA lands first; this backend keeps a distinct `flashinfer_fa2` name and module. From b41a7569a67c1bcb76c60ef6832552cc5b4f9ecb Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 02:27:50 +0800 Subject: [PATCH 10/20] docs: add Triton CUDA Graph numbers to the L20 table --- docs/accuracy.md | 7 +++++-- docs/flashinfer-fa2.md | 30 ++++++++++++++++-------------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/docs/accuracy.md b/docs/accuracy.md index 9a2b6cf6..0259caee 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -224,8 +224,11 @@ on question 182 (mean 117.1), and Triton stops that question at 40. 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: a matched 64-token fixed-work run went from 5.24/5.78/6.22 s -(eager) to 1.67/1.92/2.32 s (graphs) at concurrency 1/8/32. That is a launch- +small model for both attention backends: a matched 64-token fixed-work run at +concurrency 1/8/32 went from 5.22/5.31/5.97 s (FA2 eager) to 1.67/1.89/2.30 s +(FA2 graphs), and from 4.90/5.31/5.78 s (Triton eager) to 1.51/1.76/2.20 s +(Triton graphs). Triton graphs is 5–13% faster than FA2 graphs here, so the +gain is graph dispatch reduction rather than the FA2 kernel. That is a launch- overhead result on one L20, not a production throughput qualification; async+graphs is within noise of graphs. diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index 7aaef90e..7e467a94 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -123,20 +123,22 @@ official run used the pip `cuda-toolkit[nvcc]==13.0.3` tree. 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: - -| concurrency | Triton eager | FA2 eager | FA2 graphs | FA2 async+graphs | -|---:|---:|---:|---:|---:| -| 1 | 4.94 s (13.0 tok/s) | 5.24 s (12.2) | 1.67 s (38.4) | 1.66 s (38.7) | -| 8 | 5.50 s (93.1) | 5.78 s (88.6) | 1.92 s (267.4) | 1.88 s (272.1) | -| 32 | 6.12 s (334.5) | 6.22 s (329.4) | 2.32 s (881.7) | 2.23 s (918.8) | - -Eager FA2 is at parity with Triton; CUDA Graphs remove most per-step launch -overhead on this small model (roughly 2.6–3.1× over FA2 eager and 3.0–3.4× -over Triton eager). The official GSM8K-87 runs agree in direction (eager -784.6 s, graphs 281.5 s, with different generated lengths). Asynchronous -scheduling adds no measurable gain at these batch sizes. These are single-GPU, -single-wave diagnostics, not production throughput claims. +captures the graph. All five columns come from one session: + +| concurrency | Triton eager | Triton graphs | FA2 eager | FA2 graphs | FA2 async+graphs | +|---:|---:|---:|---:|---:|---:| +| 1 | 4.90 s (13.1 tok/s) | 1.51 s (42.3) | 5.22 s (12.3) | 1.67 s (38.4) | 1.65 s (38.8) | +| 8 | 5.31 s (96.4) | 1.76 s (290.4) | 5.31 s (96.4) | 1.89 s (270.3) | 1.86 s (275.4) | +| 32 | 5.78 s (354.3) | 2.20 s (930.0) | 5.97 s (343.1) | 2.30 s (889.6) | 2.23 s (917.5) | + +Eager FA2 is at parity with Triton eager. CUDA Graphs remove most per-step +launch overhead for **both** backends (about 3.2× over Triton eager and 2.8× +over FA2 eager on this small model); Triton graphs is then 5–13% faster than +FA2 graphs, so the graph gain is dispatch reduction, not the FA2 kernel. +The official GSM8K-87 runs agree in direction (eager 784.6 s, graphs 281.5 s, +with different generated lengths). Asynchronous scheduling adds no measurable +gain at these batch sizes. These are single-GPU, single-wave diagnostics, not +production throughput claims. 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 From aadc0a85c51828c9dd3195c2b2f36297b1d2134d Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 16:15:23 +0800 Subject: [PATCH 11/20] fix: keep FlashInfer FA2 CSR construction device-only --- vllm_rlt/kernels/flashinfer_fa2_attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm_rlt/kernels/flashinfer_fa2_attention.py b/vllm_rlt/kernels/flashinfer_fa2_attention.py index 110b423b..84e61ae6 100644 --- a/vllm_rlt/kernels/flashinfer_fa2_attention.py +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -212,7 +212,9 @@ def _write_csr(self, csr, block_tables, context_lengths): width = block_tables.shape[1] page = self.page_size pages = (context_lengths.clamp_min(0) + page - 1) // page - csr.indptr[0] = 0 + # 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) From 315a16fb444b2e20db041918b2414c4d83de9da9 Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 16:15:24 +0800 Subject: [PATCH 12/20] test: cover FA2 CSR bounds, row checks and synchronization --- tests/test_flashinfer_fa2_attention.py | 43 ++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/tests/test_flashinfer_fa2_attention.py b/tests/test_flashinfer_fa2_attention.py index c4e93342..ce852669 100644 --- a/tests/test_flashinfer_fa2_attention.py +++ b/tests/test_flashinfer_fa2_attention.py @@ -93,6 +93,21 @@ def missing(name): 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) + + @pytest.mark.gpu def test_strided_gqa_padding_and_zero_rows(): torch.manual_seed(7) @@ -176,18 +191,42 @@ 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, 145], 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() == [ - int(length) - (count - 1) * 16 for length, count in zip(lengths.tolist(), pages) + 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, width = 16, 128, 2, 64 + torch.manual_seed(17) + keys = torch.randn(32, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + values = torch.randn(32, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + q = torch.randn(rows, 16, head_dim, device="cuda", dtype=torch.bfloat16) + tables = torch.randint(0, 32, (rows, width), device="cuda", dtype=torch.int32) + lengths = torch.tensor([33, 47], device="cuda", dtype=torch.int32) + attention = FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, head_dim, page) + batch = SimpleNamespace(block_tables=tables, context_lengths=lengths) + # Warm module load and allocator caches outside the measured region. + attention.prepare_step((batch,)) + attention(q, keys, values, tables, lengths) + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + attention.prepare_step((batch,)) + finally: + torch.cuda.set_sync_debug_mode("default") + + @pytest.mark.gpu def test_graph_replay_reads_refreshed_metadata(): torch.manual_seed(13) From a355f9eaf0b2dbd2b0a08c04248ad9c3fc5e17eb Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 16:15:26 +0800 Subject: [PATCH 13/20] bench: validate --cuda-graphs before creating the output --- benchmarks/gsm8k.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index b156b2f4..6cddcf3c 100644 --- a/benchmarks/gsm8k.py +++ b/benchmarks/gsm8k.py @@ -587,12 +587,12 @@ 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") - output = Path(args.output) - output.mkdir(parents=True, exist_ok=False) attention_backend = getattr(args, "attention_backend", "triton") cuda_graphs = getattr(args, "cuda_graphs", False) if cuda_graphs and args.backend != "native": raise ValueError("--cuda-graphs applies to the native backend only") + output = Path(args.output) + output.mkdir(parents=True, exist_ok=False) metadata = { "backend": args.backend, "attention_backend": attention_backend, From 86c09b75d2f98beeb613d97ccc5134d21a7b448b Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 16:15:28 +0800 Subject: [PATCH 14/20] docs: correct FA2 graph claims and refresh post-fix measurements --- docs/accuracy.md | 13 +++++++------ docs/flashinfer-fa2.md | 38 +++++++++++++++++++++----------------- 2 files changed, 28 insertions(+), 23 deletions(-) diff --git a/docs/accuracy.md b/docs/accuracy.md index 0259caee..27ca404c 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -225,12 +225,13 @@ 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.22/5.31/5.97 s (FA2 eager) to 1.67/1.89/2.30 s -(FA2 graphs), and from 4.90/5.31/5.78 s (Triton eager) to 1.51/1.76/2.20 s -(Triton graphs). Triton graphs is 5–13% faster than FA2 graphs here, so the -gain is graph dispatch reduction rather than the FA2 kernel. That is a launch- -overhead result on one L20, not a production throughput qualification; -async+graphs is within noise of graphs. +concurrency 1/8/32 went from 5.25/5.56/6.15 s (FA2 eager) to 1.61/1.86/2.25 s +(FA2 graphs), and from 5.11/5.45/6.26 s (Triton eager) to 1.52/1.76/2.21 s +(Triton graphs). Triton graphs is 2–6% 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` free of +host synchronization (it runs it under `torch.cuda.set_sync_debug_mode`). 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 diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index 7e467a94..aeb7f8dc 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -66,7 +66,8 @@ the engine's ordinary CUDA Graph and asynchronous-scheduling options. 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, - so a step with device-owned context lengths does not read back to the host. + 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.** `__call__` copies the staging CSR into the plan's buffers; graph capture records those copies, so every replay consumes the metadata written by `prepare_step` immediately before it. The @@ -90,13 +91,14 @@ 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` (15 GPU cases): selection and +- `tests/test_flashinfer_fa2_attention.py` (18 GPU cases): 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; adapter-level CUDA Graph capture/replay - with refreshed metadata; and tiny-Ouro engine runs where eager, graphs, - async, and async+graphs produce identical token/exit-depth sequences for - LAST_EXITED and SHARED. + against a reference construction (including the empty-row last-page length); + no host synchronization in `prepare_step`; adapter-level CUDA Graph + capture/replay with refreshed metadata; and tiny-Ouro engine 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. @@ -112,7 +114,10 @@ official run used the pip `cuda-toolkit[nvcc]==13.0.3` tree. because padded graph rows and the full-width plan change the tensor-core split schedule, and each divergence changes the whole continuation. The graphs run has no length-limited runaway (mean 106.7 tokens); the eager run - hits one (question 182, mean 117.1 tokens). + hits one (question 182, mean 117.1 tokens). Both runs were repeated after + removing a host synchronization from CSR construction (a Python scalar was + assigned into a CUDA slice) and reproduced **87/87 identical sequences** + (eager 793.2 s, graphs 281.6 s). - 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. @@ -127,18 +132,17 @@ captures the graph. All five columns come from one session: | concurrency | Triton eager | Triton graphs | FA2 eager | FA2 graphs | FA2 async+graphs | |---:|---:|---:|---:|---:|---:| -| 1 | 4.90 s (13.1 tok/s) | 1.51 s (42.3) | 5.22 s (12.3) | 1.67 s (38.4) | 1.65 s (38.8) | -| 8 | 5.31 s (96.4) | 1.76 s (290.4) | 5.31 s (96.4) | 1.89 s (270.3) | 1.86 s (275.4) | -| 32 | 5.78 s (354.3) | 2.20 s (930.0) | 5.97 s (343.1) | 2.30 s (889.6) | 2.23 s (917.5) | +| 1 | 5.11 s (12.5 tok/s) | 1.52 s (42.2) | 5.25 s (12.2) | 1.61 s (39.8) | 1.59 s (40.3) | +| 8 | 5.45 s (94.0) | 1.76 s (291.2) | 5.56 s (92.1) | 1.86 s (275.4) | 1.82 s (281.7) | +| 32 | 6.26 s (327.2) | 2.21 s (927.9) | 6.15 s (333.3) | 2.25 s (908.9) | 2.17 s (943.1) | Eager FA2 is at parity with Triton eager. CUDA Graphs remove most per-step -launch overhead for **both** backends (about 3.2× over Triton eager and 2.8× -over FA2 eager on this small model); Triton graphs is then 5–13% faster than -FA2 graphs, so the graph gain is dispatch reduction, not the FA2 kernel. -The official GSM8K-87 runs agree in direction (eager 784.6 s, graphs 281.5 s, -with different generated lengths). Asynchronous scheduling adds no measurable -gain at these batch sizes. These are single-GPU, single-wave diagnostics, not -production throughput claims. +launch overhead for **both** backends (about 2.7–3.4× over eager on this small +model); Triton graphs is then 2–6% 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 793.2 s, graphs 281.6 s, with different generated lengths). These are +single-GPU, single-wave diagnostics, not production throughput claims. 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 From 4afcb44d608f87dad40caae09fd5f36b0f19e4a2 Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 19:01:01 +0800 Subject: [PATCH 15/20] feat: make FlashInfer FA2 plans deterministic and bounded --- vllm_rlt/kernels/flashinfer_fa2_attention.py | 52 +++++++++++++++++--- 1 file changed, 44 insertions(+), 8 deletions(-) diff --git a/vllm_rlt/kernels/flashinfer_fa2_attention.py b/vllm_rlt/kernels/flashinfer_fa2_attention.py index 84e61ae6..cdcb1e53 100644 --- a/vllm_rlt/kernels/flashinfer_fa2_attention.py +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -2,8 +2,16 @@ 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 created once per (rows, heads, table width) at the -cache's full table capacity, so later steps only refresh device buffers. +``__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 @@ -130,6 +138,7 @@ def __init__(self, device, dtype, head_dim, block_size): self._csr = {} self._primed = {} self._plan_misses = 0 + self._plan_creations = 0 self.info = dict( backend="flashinfer_fa2", package="flashinfer-python", @@ -137,6 +146,7 @@ def __init__(self, device, dtype, head_dim, block_size): kernel="fa2", tensor_cores=True, cuda_graph_wrapper=True, + split_kv=False, page_size=block_size, head_dim=head_dim, sm="8.9", @@ -148,6 +158,10 @@ def __init__(self, device, dtype, head_dim, block_size): 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) @@ -230,6 +244,10 @@ def _write_csr(self, csr, block_tables, context_lengths): ) 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: @@ -248,11 +266,18 @@ def _wrapper(self, rows, q_heads, kv_heads, width): 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) - # Plan for the widest table this cache can present; live CSR values are - # written before every run and unused splits are masked by the wrapper. - capacity_indptr = torch.arange(rows + 1, dtype=torch.int32, device=device) * width - capacity_indices = torch.zeros(rows * width, dtype=torch.int32, device=device) - capacity_last = 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 @@ -268,6 +293,7 @@ def _wrapper(self, rows, q_heads, kv_heads, width): 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), @@ -284,7 +310,15 @@ def _wrapper(self, rows, q_heads, kv_heads, width): 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. @@ -298,6 +332,7 @@ def _wrapper(self, rows, q_heads, kv_heads, width): last_page_len, pinned=torch.cuda.is_current_stream_capturing(), ) + self._plan_creations += 1 self._wrappers[key] = plan return plan @@ -319,8 +354,9 @@ def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): plan = self._wrapper(rows, q.shape[1], key_cache.shape[2], width) # Copied per layer on purpose: under CUDA Graph capture these nodes are # recorded and re-read the refreshed staging buffers on every replay. + # The plan's buffers cover the bucket width, so live CSR fits its prefix. plan.indptr.copy_(csr.indptr) - plan.indices.copy_(csr.indices) + plan.indices[: csr.indices.numel()].copy_(csr.indices) plan.last_page_len.copy_(csr.last_page_len) out = torch.empty_like(q) plan.wrapper.run(q, (key_cache, value_cache), out=out) From 55add980ebe4a98a04a1aadca90b04a63d54388b Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 19:01:02 +0800 Subject: [PATCH 16/20] feat: expose the resolved FA2 kernel through selection info --- vllm_rlt/attention/backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 1c0fb021..a7c3147b 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -94,7 +94,9 @@ def _create_explicit(backend, device, dtype, head_dim, block_size): from vllm_rlt.kernels.flashinfer_fa2_attention import FlashInferFA2PagedAttention kernel = FlashInferFA2PagedAttention(device, dtype, head_dim, block_size) - return AttentionBackend(kernel, _FLASHINFER_FA2_CAPABILITIES, dict(kernel.info)) + # 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)) From 7123886b3d1b4d070539d4bcc1702e60ea64bf76 Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 19:01:03 +0800 Subject: [PATCH 17/20] test: cover FA2 plan bounds, split invariance and wide-table graphs --- tests/test_flashinfer_fa2_attention.py | 154 +++++++++++++++++++++++-- 1 file changed, 145 insertions(+), 9 deletions(-) diff --git a/tests/test_flashinfer_fa2_attention.py b/tests/test_flashinfer_fa2_attention.py index ce852669..113cbe8f 100644 --- a/tests/test_flashinfer_fa2_attention.py +++ b/tests/test_flashinfer_fa2_attention.py @@ -1,5 +1,6 @@ """FlashInfer FA2 adapter: SM89 selection, CSR metadata and eager decode.""" +import dataclasses import sys import types from types import SimpleNamespace @@ -108,6 +109,70 @@ def test_wrapper_rejects_unsupported_head_ratio(fake_flashinfer): 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) @@ -207,24 +272,31 @@ def test_vectorized_csr_matches_reference(): @pytest.mark.gpu def test_prepare_step_does_not_synchronize(): - page, head_dim, rows, width = 16, 128, 2, 64 + page, head_dim, rows = 16, 128, 2 torch.manual_seed(17) - keys = torch.randn(32, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] - values = torch.randn(32, 3, page, 16, head_dim, device="cuda", dtype=torch.bfloat16)[:, 1] + 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) - tables = torch.randint(0, 32, (rows, width), device="cuda", dtype=torch.int32) 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) - batch = SimpleNamespace(block_tables=tables, context_lengths=lengths) - # Warm module load and allocator caches outside the measured region. - attention.prepare_step((batch,)) - attention(q, keys, values, tables, lengths) + # 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: - attention.prepare_step((batch,)) + 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 @@ -282,6 +354,57 @@ def test_paged_attention_matches_triton(q_heads, kv_heads): 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( @@ -334,3 +457,16 @@ def test_tiny_ouro_execution_variants_match_eager(layout, execution): 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 From d522a0d8ff6b1e2dc4d20806a3e5d468c28dc4d9 Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 19:01:05 +0800 Subject: [PATCH 18/20] docs: document deterministic FA2 plans and refresh measurements --- docs/accuracy.md | 28 ++++++------ docs/flashinfer-fa2.md | 97 ++++++++++++++++++++++++++++-------------- 2 files changed, 81 insertions(+), 44 deletions(-) diff --git a/docs/accuracy.md b/docs/accuracy.md index 27ca404c..65cbb8ad 100644 --- a/docs/accuracy.md +++ b/docs/accuracy.md @@ -211,27 +211,31 @@ On one L20 (SM89, torch 2.13.0, CUDA 13.0, FlashInfer 0.7.2.post1, Triton | run | correct | accuracy | gate | |---|---:|---:|---| | triton | 61/87 | 70.11% | pass (+2.30 pp) | -| flashinfer_fa2 eager | 60/87 | 68.97% | pass (+1.15 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 share 69/87 generated token sequences. Padded graph rows -and the full-width plan change the tensor-core split schedule, so 18 first -divergences are BF16 ties (logits maxabs <= 0.17, top1/top2 margins 0–0.125) -and each one changes the whole continuation. The graphs run has no -length-limited runaway (mean 106.7 tokens); the eager run reaches 1024 tokens -on question 182 (mean 117.1), and Triton stops that question at 40. +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.25/5.56/6.15 s (FA2 eager) to 1.61/1.86/2.25 s -(FA2 graphs), and from 5.11/5.45/6.26 s (Triton eager) to 1.52/1.76/2.21 s -(Triton graphs). Triton graphs is 2–6% faster than FA2 graphs here, so the gain +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` free of -host synchronization (it runs it under `torch.cuda.set_sync_debug_mode`). +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 diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index aeb7f8dc..d4b53cfc 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -53,15 +53,29 @@ the engine's ordinary CUDA Graph and asynchronous-scheduling options. 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 full table capacity but takes the schedule - from the live CSR buffers and masks unused splits, which is the graph - contract FlashInfer documents for this wrapper. -- **Plan cache.** One plan per `(rows, q_heads, kv_heads, table_width)`, created - on first use with a synthetic full-capacity table. Each plan owns its CSR - buffers and a workspace sized by `workspace_size(...)`; FlashInfer keeps - plan-time tables in that workspace, so plans cannot share it. Plans captured - by a CUDA Graph are pinned; eager plans are evicted first when the cache - (32 entries) is full. Every plan is created outside graph capture. + 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 @@ -91,14 +105,15 @@ 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` (18 GPU cases): selection and - rejection paths; strided/GQA/MHA KV against both Torch and Triton kernels; +- `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`; adapter-level CUDA Graph - capture/replay with refreshed metadata; and tiny-Ouro engine runs where - eager, graphs, async, and async+graphs produce identical token/exit-depth - sequences for LAST_EXITED and SHARED. + 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. @@ -107,17 +122,17 @@ official run used the pip `cuda-toolkit[nvcc]==13.0.3` tree. | run | correct | accuracy | gate | |---|---:|---:|---| | Triton eager | 61/87 | 70.11% | pass | -| FA2 eager (tensor cores) | 60/87 | 68.97% | pass (+1.15 pp) | +| 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 share **69/87** token sequences; 18 diverge at a BF16 tie - because padded graph rows and the full-width plan change the tensor-core - split schedule, and each divergence changes the whole continuation. The - graphs run has no length-limited runaway (mean 106.7 tokens); the eager run - hits one (question 182, mean 117.1 tokens). Both runs were repeated after - removing a host synchronization from CSR construction (a Python scalar was - assigned into a CUDA slice) and reproduced **87/87 identical sequences** - (eager 793.2 s, graphs 281.6 s). + 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. @@ -132,18 +147,32 @@ captures the graph. All five columns come from one session: | concurrency | Triton eager | Triton graphs | FA2 eager | FA2 graphs | FA2 async+graphs | |---:|---:|---:|---:|---:|---:| -| 1 | 5.11 s (12.5 tok/s) | 1.52 s (42.2) | 5.25 s (12.2) | 1.61 s (39.8) | 1.59 s (40.3) | -| 8 | 5.45 s (94.0) | 1.76 s (291.2) | 5.56 s (92.1) | 1.86 s (275.4) | 1.82 s (281.7) | -| 32 | 6.26 s (327.2) | 2.21 s (927.9) | 6.15 s (333.3) | 2.25 s (908.9) | 2.17 s (943.1) | +| 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 (about 2.7–3.4× over eager on this small -model); Triton graphs is then 2–6% faster than FA2 graphs, so the graph gain is +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 793.2 s, graphs 281.6 s, with different generated lengths). These are +(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. @@ -154,7 +183,11 @@ used here and are not reproduced. - 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. -- Size plan workspaces per graph-entry width; full-capacity plans allocate the - widest table even for short contexts. +- Write the live CSR directly into the plan buffers at prime time (currently + three per-layer copies, about 0.5 MiB per layer at 32 rows and 4096 pages). +- 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. +- Add a CPU test for the GSM8K `--cuda-graphs` / `--attention-backend` flags. - Reconcile with #112 if XQA lands first; this backend keeps a distinct `flashinfer_fa2` name and module. From 9e170a442a7a10b1457e4a257f49ae553152f785 Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 20:03:50 +0800 Subject: [PATCH 19/20] perf: copy FA2 CSR metadata once per traversal --- docs/flashinfer-fa2.md | 16 ++++----- vllm_rlt/kernels/flashinfer_fa2_attention.py | 34 +++++++++++++++----- 2 files changed, 34 insertions(+), 16 deletions(-) diff --git a/docs/flashinfer-fa2.md b/docs/flashinfer-fa2.md index d4b53cfc..19eee527 100644 --- a/docs/flashinfer-fa2.md +++ b/docs/flashinfer-fa2.md @@ -82,11 +82,11 @@ the engine's ordinary CUDA Graph and asynchronous-scheduling options. (`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.** `__call__` copies the staging CSR into - the plan's buffers; graph capture records those copies, 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. +- **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 @@ -183,11 +183,11 @@ used here and are not reproduced. - 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 (currently - three per-layer copies, about 0.5 MiB per layer at 32 rows and 4096 pages). +- 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. -- Add a CPU test for the GSM8K `--cuda-graphs` / `--attention-backend` flags. - Reconcile with #112 if XQA lands first; this backend keeps a distinct `flashinfer_fa2` name and module. diff --git a/vllm_rlt/kernels/flashinfer_fa2_attention.py b/vllm_rlt/kernels/flashinfer_fa2_attention.py index cdcb1e53..4269c5c6 100644 --- a/vllm_rlt/kernels/flashinfer_fa2_attention.py +++ b/vllm_rlt/kernels/flashinfer_fa2_attention.py @@ -68,7 +68,15 @@ class _CSR: and the raw pointers recorded by its capture valid. """ - __slots__ = ("indptr", "indices", "last_page_len", "block_tables", "context_lengths") + __slots__ = ( + "indptr", + "indices", + "last_page_len", + "block_tables", + "context_lengths", + "synced", + "captured", + ) def __init__(self, block_tables, context_lengths): rows, width = block_tables.shape @@ -78,6 +86,10 @@ def __init__(self, block_tables, context_lengths): 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 @@ -87,7 +99,8 @@ 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 ``__call__`` copies it in, so replays see new metadata. + 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") @@ -184,6 +197,7 @@ def prepare_step(self, batches): 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 @@ -352,12 +366,16 @@ def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): 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) - # Copied per layer on purpose: under CUDA Graph capture these nodes are - # recorded and re-read the refreshed staging buffers on every replay. - # The plan's buffers cover the bucket width, so live CSR fits its prefix. - plan.indptr.copy_(csr.indptr) - plan.indices[: csr.indices.numel()].copy_(csr.indices) - plan.last_page_len.copy_(csr.last_page_len) + # 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)) From 765fe415a029bc38aa8aec32b9808ae3dfa6c73f Mon Sep 17 00:00:00 2001 From: bohaotu Date: Mon, 2 Nov 2026 20:03:52 +0800 Subject: [PATCH 20/20] bench: fail fast for non-native --cuda-graphs --- benchmarks/gsm8k.py | 4 ++-- tests/test_gsm8k_accuracy.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index 6cddcf3c..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 @@ -589,8 +591,6 @@ def run(args): raise RuntimeError("Run with one scheduler-assigned GPU") attention_backend = getattr(args, "attention_backend", "triton") cuda_graphs = getattr(args, "cuda_graphs", False) - if cuda_graphs and args.backend != "native": - raise ValueError("--cuda-graphs applies to the native backend only") output = Path(args.output) output.mkdir(parents=True, exist_ok=False) metadata = { 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()