diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c6afe035f..26c0ea0ebf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - Add Jupyter notebook tutorial for using amd-flashinfer on ROCm (#213) @diptorupd - Add ROCm profiler module (`flashinfer/rocm_profiler/`) and FA2 single-prefill benchmark driver using rocprofv3 (#205) @diptorupd - Gate `torch.compile` integration behind `FLASHINFER_USE_TORCH_CUSTOM_OPS`, with HIP pytest coverage (#210) @demandal25 +- Support CUDA-graph capture on the AITER batch-decode path via an explicit `backend="aiter"`: capture at a maximum sequence length and replay for shorter sequences (the launch grid and `.so` variant are fixed at capture-time shapes; the kernel early-exits per sequence on `context_lens`). Previously `backend="aiter"` raised under `use_cuda_graph=True`. `backend="auto"` continues to use the in-tree `fa2` kernel under capture (fa2's graph path is capacity-based and capture-order-independent). Adds `benchmarks/rocm_benchmarks/bench_decode_graph_hip.py` and HIP graph-replay tests. @demandal25 ## Changed diff --git a/README.md b/README.md index 5b06a92117..1efa013787 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ kernel for non-attention ops). **AITER** = ROCm AITER backend. | Kernel | HIP | AITER | `backend="auto"` resolves to | Notes | | :--- | :---: | :---: | :--- | :--- | | **Single decode attention** | ✅ `fa2` | — | HIP | MHA / GQA / MQA | -| **Batch decode attention (paged)** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + `pos_encoding_mode="NONE"` + no CUDA-graph + `use_tensor_cores=False`; else **HIP** | MHA / GQA / MQA; **fp8 KV-cache (E4M3FNUZ)** on the HIP path; sliding-window on the AITER path; CUDA-graph auto-routes back to HIP | +| **Batch decode attention (paged)** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + `pos_encoding_mode="NONE"` + no CUDA-graph + `use_tensor_cores=False`; else **HIP** | MHA / GQA / MQA; **fp8 KV-cache (E4M3FNUZ)** on the HIP path; sliding-window on the AITER path; CUDA-graph capture on the AITER path is **opt-in via `backend="aiter"`** (grid + `.so` fixed at capture-time shapes — capture at max seq len); `backend="auto"` uses HIP `fa2` under capture | | **Single prefill attention** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + no custom mask + equal Q/KV dtypes & head dims + `pos_encoding_mode="NONE"`; else **HIP** | MHA / GQA / MQA; fp8 WIP | | **Batch prefill attention (paged + ragged)** | ✅ `fa2` | ✅ | Same auto criteria as single prefill | MHA / GQA / MQA; fp8 WIP. AITER native page sizes: `{16, 1024}` (`{128, 256, 1024}` on `amd-aiter==0.1.10`); other sizes go through a gather on the AITER path | | **Cascade attention** | ✅ | — | HIP | Two-level shared-prefix attention; a fused single-kernel HIP variant is gated behind `FLASHINFER_HIP_FUSED_CASCADE=1` | @@ -354,9 +354,14 @@ Backend-specific exceptions to "auto picks AITER when supported": * `rmsnorm`: `backend="auto"` picks the AITER C++ path (CK `rmsnorm2d`) only for 2-D fp16/bf16 inputs whose weight dtype matches; 3-D inputs, fp32, or a mismatched weight dtype fall back to the HIP `native` kernel. -* `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True` - force `auto` back to `fa2` (AITER decode does not support either), - and `pos_encoding_mode != "NONE"` raises under `backend="aiter"`. +* `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True` force `auto` + back to `fa2`, and `pos_encoding_mode != "NONE"` raises under + `backend="aiter"`. CUDA-graph capture on the AITER decode path is available + via an explicit `backend="aiter"` (not `auto`): capture at your maximum + sequence length — the grid and `.so` variant are fixed at capture-time shapes + and the kernel early-exits per sequence, so replays *shorter* than captured + are correct but *longer* ones are not. fa2's graph path is capacity-based and + carries no such constraint. Unless you are using the prebuilt Docker image, install AITER separately via one of the options below. @@ -396,7 +401,7 @@ you need any of them: `backend="fa2"` for attention wrappers, or * `q_dtype != kv_dtype` (mixed-precision Q/KV is unsupported) * `head_dim_qk != head_dim_vo` (e.g. DeepSeek-style MLA with 192/128 head dims) * `pos_encoding_mode != "NONE"` (AITER attention paths only support `"NONE"`) -* batch decode: `use_cuda_graph=True` or `use_tensor_cores=True` +* batch decode: `use_tensor_cores=True` * the `aiter` Python package is not importable **Features silently ignored on the AITER path** (kwargs are accepted by diff --git a/benchmarks/rocm_benchmarks/bench_batch_decode_hip.py b/benchmarks/rocm_benchmarks/bench_batch_decode_hip.py new file mode 100644 index 0000000000..867ffb0ade --- /dev/null +++ b/benchmarks/rocm_benchmarks/bench_batch_decode_hip.py @@ -0,0 +1,205 @@ +""" +Copyright (c) 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Batch paged-decode benchmark: fa2 (in-tree HIP) vs aiter (PA v1), eager. + +Baseline for the "HIP-graph-safe fast decode" work item +(docs/rocm_library_optimization_plan.md §3 #1): today AITER decode falls back +to fa2 under CUDA-graph capture, so the fast path never runs in the serving +hot loop. This bench measures the eager fa2-vs-aiter gap (the win we forfeit +under graphs) across a serving-shaped batch x kv_len sweep. Once the +capacity-sized-grid change lands, add a --graph path here to compare +aiter-under-graph vs fa2-under-graph and to measure the idle-block +over-provisioning cost (the one open perf item in the plan). + +Shapes: decode (q_len=1), GQA 32/8, HD=128, bf16, page_size=16. +Sweep: batch x kv_len. + +Run: + python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py # full pipeline + python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --timing-only # no profiling + python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --backend fa2 + python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --backend aiter + python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --counters stall + +Design note: bench flags are parsed at module level because rocprofv3 +re-executes this script as a subprocess per PMC pass with the same sys.argv. +Module-level parsing ensures the subprocess builds identical configs to the +outer timing run. +""" + +import argparse +import logging +import sys +from pathlib import Path + +import torch + +import flashinfer +from flashinfer.jit.core import logger as _jit_logger + +# Suppress routine JIT INFO/DEBUG output; WARNING still surfaces compile errors. +_jit_logger.setLevel(logging.WARNING) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "rocm_profiler")) +from rocm_profiler import KernelConfig, RocmProfiler + +# --------------------------------------------------------------------------- +# Bench-script-level argument parsing (see design note above) +# --------------------------------------------------------------------------- +_bench_parser = argparse.ArgumentParser(add_help=False) +_bench_parser.add_argument( + "--counters", + default="roofline", + metavar="PRESET_OR_FILE", + help=( + "Counter preset name ('roofline', 'occupancy', 'stall', 'compute', " + "'memory', 'basic') or path to a rocprofv3 YAML file. Default: roofline." + ), +) +_bench_parser.add_argument( + "--label", + default=None, + metavar="PREFIX", + help="Output-file label prefix (default: 'decode' for roofline, 'decode_' otherwise).", +) +_bench_parser.add_argument( + "--backend", + default="both", + choices=["fa2", "aiter", "both"], + help="Which decode backend(s) to sweep. Default: both.", +) +_bench_args, _ = _bench_parser.parse_known_args() + +_counters = _bench_args.counters +_label = ( + _bench_args.label + if _bench_args.label is not None + else ("decode" if _counters == "roofline" else f"decode_{_counters}") +) + +# --------------------------------------------------------------------------- +# Sweep configuration — decode is q_len=1; shapes ~ Llama-70B TP1 decode. +# --------------------------------------------------------------------------- +_NUM_QO_HEADS = 32 +_NUM_KV_HEADS = 8 +_HEAD_DIM = 128 +_DTYPE = torch.bfloat16 +_PAGE_SIZE = 16 +_BATCHES = [1, 8, 32, 128, 256] +_KV_LENS = [1024, 2048, 4096, 8192] + +_OUTPUT_DIR = str(Path(__file__).parent) + + +def _flops(kv_len: int, num_qo_heads: int, head_dim: int) -> int: + # decode: q_len=1 → attended = kv_len; QK^T + PV ≈ 2 matmuls → factor 4. + return kv_len * num_qo_heads * head_dim * 4 + + +def _bytes(kv_len: int, num_qo_heads: int, num_kv_heads: int, head_dim: int) -> int: + # Dominated by reading K and V (q_len=1). 2 bytes/elem (bf16). + return 2 * head_dim * (2 * num_qo_heads + 2 * kv_len * num_kv_heads) + + +def _build_paged_kv(batch, kv_len, page_size, num_kv_heads, head_dim, dtype, device): + """Build batch-paged decode KV + query tensors for one (batch, kv_len).""" + num_full_pages, last_tokens = divmod(kv_len, page_size) + if last_tokens == 0: + last_tokens = page_size + else: + num_full_pages += 1 + total_pages = num_full_pages * batch + + kv_data = torch.randn( + total_pages, 2, page_size, num_kv_heads, head_dim, dtype=dtype, device=device + ) + kv_indptr = ( + torch.arange(batch + 1, dtype=torch.int32, device=device) * num_full_pages + ) + _rng = torch.Generator(device=device).manual_seed(42) + kv_indices = torch.randperm( + total_pages, dtype=torch.int32, device=device, generator=_rng + ) + kv_last_page_len = torch.full( + (batch,), last_tokens, dtype=torch.int32, device=device + ) + q = torch.randn(batch, _NUM_QO_HEADS, head_dim, dtype=dtype, device=device) + return q, kv_data, kv_indptr, kv_indices, kv_last_page_len + + +def _make_wrapper_config(backend, batch, kv_len): + device = torch.device("cuda") + q, kv_data, kv_indptr, kv_indices, kv_last_page_len = _build_paged_kv( + batch, kv_len, _PAGE_SIZE, _NUM_KV_HEADS, _HEAD_DIM, _DTYPE, device + ) + ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device) + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(ws, "NHD", backend=backend) + wrapper.plan( + kv_indptr, + kv_indices, + kv_last_page_len, + _NUM_QO_HEADS, + _NUM_KV_HEADS, + _HEAD_DIM, + _PAGE_SIZE, + pos_encoding_mode="NONE", + q_data_type=_DTYPE, + kv_data_type=_DTYPE, + ) + return KernelConfig( + name=f"{backend}_b{batch}_kv{kv_len}", + run_fn=torch.inference_mode()(lambda q=q, kv=kv_data, w=wrapper: w.run(q, kv)), + theoretical_flops=batch * _flops(kv_len, _NUM_QO_HEADS, _HEAD_DIM), + theoretical_bytes=batch + * _bytes(kv_len, _NUM_QO_HEADS, _NUM_KV_HEADS, _HEAD_DIM), + num_tokens=batch, # one decoded token per sequence + label=f"{backend:<5s} b={batch:>3d} kv={kv_len:>5d}", + ) + + +@torch.inference_mode() +def _make_configs() -> list[KernelConfig]: + backends = ( + ["fa2", "aiter"] if _bench_args.backend == "both" else [_bench_args.backend] + ) + if "aiter" in backends and not flashinfer.aiter_utils.is_aiter_supported( + torch.device("cuda:0") + ): + print("AITER not supported on this device; dropping the aiter sweep.") + backends = [b for b in backends if b != "aiter"] + + configs: list[KernelConfig] = [] + for backend in backends: + for batch in _BATCHES: + for kv_len in _KV_LENS: + configs.append(_make_wrapper_config(backend, batch, kv_len)) + return configs + + +if __name__ == "__main__": + _skip_gpu = "--replot" in sys.argv or "--list-presets" in sys.argv + profiler = RocmProfiler( + configs=[] if _skip_gpu else _make_configs(), + num_warmup=3, + dry_run_ms=100, + repeat_ms=1000, + counters=_counters, + kernel_name_regex="", + output_dir=_OUTPUT_DIR, + label=_label, + roofline=(_counters == "roofline"), + ) + profiler.run() diff --git a/benchmarks/rocm_benchmarks/bench_decode_graph_hip.py b/benchmarks/rocm_benchmarks/bench_decode_graph_hip.py new file mode 100644 index 0000000000..b122390075 --- /dev/null +++ b/benchmarks/rocm_benchmarks/bench_decode_graph_hip.py @@ -0,0 +1,210 @@ +""" +Copyright (c) 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Acceptance test + timing for graph-capturable AITER decode (plan §3 #1). + +Verifies that AITER PA v1 decode, captured once at a maximum sequence length, +replays *correctly* for shorter sequences (the kernel early-exits per-seq on +context_lens), and times aiter-under-graph vs fa2-under-graph. + +Run: + python benchmarks/rocm_benchmarks/bench_decode_graph_hip.py +""" + +import os +import sys +import time + +import torch + +import flashinfer + +# Fixed problem geometry (Llama-70B TP1-ish decode). BATCH/CAP_SEQ overridable +# via env for quick scaling checks: FI_GRAPH_BATCH, FI_GRAPH_CAP_SEQ. +BATCH = int(os.environ.get("FI_GRAPH_BATCH", "16")) +PAGE = 16 +NUM_QO, NUM_KV, HD = 32, 8, 128 +DTYPE = torch.bfloat16 +CAP_SEQ = int( + os.environ.get("FI_GRAPH_CAP_SEQ", "4096") +) # capacity captured; replays <= this +CAP_PAGES_PER_SEQ = (CAP_SEQ + PAGE - 1) // PAGE +TOTAL_PAGES = BATCH * CAP_PAGES_PER_SEQ +DEVICE = torch.device("cuda") + + +def _layout_for(seq_len: int): + """Uniform per-seq kv_len=seq_len over the fixed page pool. + + Each sequence keeps a stable reserved block of CAP_PAGES_PER_SEQ pages; a + shorter seq_len uses only the first `npages` of its block. This mirrors a + real fixed-capacity paged-KV pool (stable per-seq page mapping) rather than + repacking pages when seq_len < CAP_SEQ. + """ + npages = (seq_len + PAGE - 1) // PAGE + last = seq_len - (npages - 1) * PAGE + indptr = torch.arange(BATCH + 1, dtype=torch.int32, device=DEVICE) * npages + base = (torch.arange(BATCH, device=DEVICE) * CAP_PAGES_PER_SEQ).view(-1, 1) + offs = torch.arange(npages, device=DEVICE).view(1, -1) + indices = (base + offs).reshape(-1).to(torch.int32) + last_page = torch.full((BATCH,), last, dtype=torch.int32, device=DEVICE) + return indptr, indices, last_page + + +def _reference(q, kv, seq_len, backend): + """Eager (non-graph) wrapper result for the same inputs.""" + ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=DEVICE) + w = flashinfer.BatchDecodeWithPagedKVCacheWrapper(ws, "NHD", backend=backend) + indptr, indices, last_page = _layout_for(seq_len) + w.plan( + indptr, + indices, + last_page, + NUM_QO, + NUM_KV, + HD, + PAGE, + pos_encoding_mode="NONE", + q_data_type=DTYPE, + kv_data_type=DTYPE, + ) + return w.run(q, kv) + + +def _make_graph_wrapper(backend): + ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=DEVICE) + indptr_buf = torch.empty(BATCH + 1, dtype=torch.int32, device=DEVICE) + indices_buf = torch.empty(TOTAL_PAGES, dtype=torch.int32, device=DEVICE) + last_page_buf = torch.empty(BATCH, dtype=torch.int32, device=DEVICE) + w = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + ws, + "NHD", + use_cuda_graph=True, + backend=backend, + paged_kv_indptr_buffer=indptr_buf, + paged_kv_indices_buffer=indices_buf, + paged_kv_last_page_len_buffer=last_page_buf, + ) + return w + + +def _capture(w, q_static, kv): + """plan() at capacity, then capture run() into a static output.""" + indptr, indices, last_page = _layout_for(CAP_SEQ) + w.plan( + indptr, + indices, + last_page, + NUM_QO, + NUM_KV, + HD, + PAGE, + pos_encoding_mode="NONE", + q_data_type=DTYPE, + kv_data_type=DTYPE, + ) + # warmup (also triggers dlopen / first alloc outside capture) + for _ in range(3): + w.run(q_static, kv) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_static = w.run(q_static, kv) + return g, out_static + + +def main(): + if not flashinfer.aiter_utils.is_aiter_supported(DEVICE): + print("AITER not supported on this device. Exiting.") + return + + torch.manual_seed(0) + kv = torch.randn(TOTAL_PAGES, 2, PAGE, NUM_KV, HD, dtype=DTYPE, device=DEVICE) + q_static = torch.randn(BATCH, NUM_QO, HD, dtype=DTYPE, device=DEVICE) + + print(f"Backend resolves under graph capture (batch={BATCH}, cap_seq={CAP_SEQ}):") + w = _make_graph_wrapper("aiter") + g, out_static = _capture(w, q_static, kv) + print(f" captured backend = {w._backend!r}") + assert w._backend == "aiter", f"expected aiter under graph, got {w._backend!r}" + + print("\nReplay correctness across seq_len <= capacity (vs eager aiter):") + ok = True + for seq_len in [512, 1024, 2048, 4096]: + # fresh q contents for this step + q_new = torch.randn(BATCH, NUM_QO, HD, dtype=DTYPE, device=DEVICE) + q_static.copy_(q_new) + # update fixed buffers to the real (shorter) layout — NOT captured + indptr, indices, last_page = _layout_for(seq_len) + w.plan( + indptr, + indices, + last_page, + NUM_QO, + NUM_KV, + HD, + PAGE, + pos_encoding_mode="NONE", + q_data_type=DTYPE, + kv_data_type=DTYPE, + ) + g.replay() + torch.cuda.synchronize() + ref = _reference(q_new, kv, seq_len, "aiter") + max_diff = (out_static.float() - ref.float()).abs().max().item() + good = torch.allclose(out_static, ref, atol=2e-2, rtol=2e-2) + ok = ok and good + print( + f" seq_len={seq_len:>5d} max|graph-eager|={max_diff:.4f} " + f"{'PASS' if good else 'FAIL'}" + ) + + print(f"\nOverall correctness: {'PASS' if ok else 'FAIL'}") + + # ── timing: aiter-under-graph vs fa2-under-graph (replay only) ────────── + print(f"\nUnder-graph replay latency (seq_len=4096, batch={BATCH}):") + for backend in ["fa2", "aiter"]: + wl = _make_graph_wrapper(backend) + gl, _ = _capture(wl, q_static, kv) + # set a mid-size real layout + indptr, indices, last_page = _layout_for(4096) + wl.plan( + indptr, + indices, + last_page, + NUM_QO, + NUM_KV, + HD, + PAGE, + pos_encoding_mode="NONE", + q_data_type=DTYPE, + kv_data_type=DTYPE, + ) + for _ in range(5): + gl.replay() + torch.cuda.synchronize() + n = 200 + t0 = time.perf_counter() + for _ in range(n): + gl.replay() + torch.cuda.synchronize() + ms = (time.perf_counter() - t0) / n * 1e3 + print(f" {backend:<5s} {ms:.4f} ms / replay") + + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main() or 0) diff --git a/flashinfer/decode_rocm.py b/flashinfer/decode_rocm.py index 25c492fd65..89d0caae68 100644 --- a/flashinfer/decode_rocm.py +++ b/flashinfer/decode_rocm.py @@ -394,6 +394,10 @@ def _fake_run_batch_decode_aiter( # wrapper will be dispatched through the FA2 shadow plan (AITER PA v1 does not output LSE). _aiter_lse_fallback_warned: set[torch.device] = set() +# One-time-per-device warning about AITER decode's capture-at-max-seq-len contract +# under CUDA-graph capture (opt-in via explicit backend="aiter"). +_aiter_graph_capture_warned: set[torch.device] = set() + def _aiter_pa_v1_resolve( *, @@ -1125,10 +1129,13 @@ def plan( kv_lens_arr_host = seq_lens.cpu() # Resolve auto → concrete backend. AITER decode requires use_tensor_cores=False - # (the AITER PA v1 kernel handles its own dispatch internally). CUDA-graph - # capture is excluded: AITER's launch grid is sized from per-plan scalars - # (max_kv_len, max_blocks_per_seq) that get baked into the captured graph - # and cannot be widened on replay without re-capturing. + # (the AITER PA v1 kernel handles its own dispatch internally). Under CUDA-graph + # capture, `auto` stays on fa2: fa2's graph path is capacity-based and correct + # regardless of capture-vs-replay sizes, whereas AITER's launch grid and .so + # variant are fixed at the shapes seen when the graph is captured and require + # capturing at the maximum sequence length. AITER decode under CUDA graph is + # therefore opt-in via an explicit backend="aiter" (see the capture-at-max + # warning below), not something `auto` selects silently. if self._backend == "auto": if self.use_tensor_cores or self.is_cuda_graph_enabled: self._backend = "fa2" @@ -1155,14 +1162,18 @@ def plan( f"AITER decode backend requires pos_encoding_mode='NONE', " f"got {pos_encoding_mode!r}" ) - if self.is_cuda_graph_enabled: - raise ValueError( - "AITER decode backend is incompatible with CUDA-graph capture: " - "the kernel's launch grid is sized from per-plan scalars " - "(max_kv_len, max_blocks_per_seq) that are baked into the " - "captured graph at capture time. Use backend='fa2' for " - "CUDA-graph workflows, or backend='auto' which routes around " - "this automatically." + if ( + self.is_cuda_graph_enabled + and self.device not in _aiter_graph_capture_warned + ): + _aiter_graph_capture_warned.add(self.device) + logger.warning( + "AITER decode under CUDA-graph capture: the launch grid and kernel " + "variant are fixed at the shapes seen when the graph is captured " + "(unlike fa2, whose graph path is capacity-based). Capture at your " + "maximum sequence length — replays with sequences longer than " + "captured will be incorrect. Use backend='fa2' if you need " + "capture-order-independent CUDA-graph decode." ) self._max_kv_len = int(max(kv_lens_arr_host).item()) # max blocks per seq across the batch — needed to size the dense block_tables. diff --git a/tests/rocm_tests/test_batch_decode_aiter_hip.py b/tests/rocm_tests/test_batch_decode_aiter_hip.py index 63def3f2c7..9e3ceee411 100644 --- a/tests/rocm_tests/test_batch_decode_aiter_hip.py +++ b/tests/rocm_tests/test_batch_decode_aiter_hip.py @@ -181,34 +181,10 @@ def test_batch_decode_aiter_rejects_invalid_config(): kv_data_type=torch.float16, ) - # CUDA-graph capture not supported with explicit backend="aiter". - indptr_buf = torch.empty(2, dtype=torch.int32, device=device) - indices_buf = torch.empty(8, dtype=torch.int32, device=device) - last_page_len_buf = torch.empty(1, dtype=torch.int32, device=device) - w = flashinfer.BatchDecodeWithPagedKVCacheWrapper( - workspace, - "NHD", - use_cuda_graph=True, - paged_kv_indptr_buffer=indptr_buf, - paged_kv_indices_buffer=indices_buf, - paged_kv_last_page_len_buffer=last_page_len_buf, - backend="aiter", - ) - indptr_buf.copy_(indptr) - indices_buf[:1].copy_(indices) - last_page_len_buf.copy_(last_page_len) - with pytest.raises(ValueError, match="CUDA-graph"): - w.plan( - indptr_buf, - indices_buf[:1], - last_page_len_buf, - 8, - 8, - 128, - 16, - q_data_type=torch.float16, - kv_data_type=torch.float16, - ) + # NOTE: CUDA-graph capture with backend="aiter" is now supported (the grid + # and .so variant are fixed at capture-time shapes; the kernel early-exits + # per-seq on context_lens). Positive coverage lives in + # test_batch_decode_aiter_cuda_graph_replay below. @requires_aiter @@ -357,8 +333,10 @@ def test_batch_decode_aiter_return_lse_via_fa2(dtype, window_left): @requires_aiter def test_batch_decode_auto_routes_cuda_graph_to_fa2(): - """backend='auto' with use_cuda_graph=True must route to fa2 (AITER doesn't - support graph capture).""" + """backend='auto' with use_cuda_graph=True routes to fa2. fa2's graph path is + capacity-based (correct regardless of capture-vs-replay sizes); AITER decode + under graph is opt-in via backend='aiter' (capture-at-max contract), so auto + does not select it silently.""" device = torch.device("cuda:0") workspace = torch.zeros(8 * 1024 * 1024, dtype=torch.uint8, device=device) indptr_buf = torch.empty(2, dtype=torch.int32, device=device) @@ -389,3 +367,108 @@ def test_batch_decode_auto_routes_cuda_graph_to_fa2(): kv_data_type=torch.float16, ) assert w._backend == "fa2" + + +@requires_aiter +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_batch_decode_aiter_cuda_graph_replay(dtype): + """Opt-in AITER decode under CUDA graph (explicit backend='aiter'): capture + once at a maximum sequence length, then replay for a shorter sequence; the + result must match an eager AITER run (the kernel early-exits per-seq on + context_lens). Captures at cap_seq and replays shorter — the supported + capture-at-max usage.""" + device = torch.device("cuda:0") + batch, page, num_qo, num_kv, hd = 4, 16, 8, 8, 128 + cap_seq, replay_seq = 2048, 512 + cap_pages = (cap_seq + page - 1) // page + total_pages = batch * cap_pages + + kv = torch.randn(total_pages, 2, page, num_kv, hd, dtype=dtype, device=device) + q = torch.randn(batch, num_qo, hd, dtype=dtype, device=device) + + def layout(seq_len): + npages = (seq_len + page - 1) // page + last = seq_len - (npages - 1) * page + indptr = torch.arange(batch + 1, dtype=torch.int32, device=device) * npages + # Each sequence keeps a stable reserved block of cap_pages in the fixed + # pool; a shorter seq_len just uses the first `npages` of its block. This + # models real paged-KV (stable per-seq page pool) and exercises the + # capture-at-max contract faithfully. + base = (torch.arange(batch, device=device) * cap_pages).view(-1, 1) + offs = torch.arange(npages, device=device).view(1, -1) + indices = (base + offs).reshape(-1).to(torch.int32) + last_page = torch.full((batch,), last, dtype=torch.int32, device=device) + return indptr, indices, last_page + + ws = torch.zeros(64 * 1024 * 1024, dtype=torch.uint8, device=device) + indptr_buf = torch.empty(batch + 1, dtype=torch.int32, device=device) + indices_buf = torch.empty(total_pages, dtype=torch.int32, device=device) + last_page_buf = torch.empty(batch, dtype=torch.int32, device=device) + w = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + ws, + "NHD", + use_cuda_graph=True, + backend="aiter", + paged_kv_indptr_buffer=indptr_buf, + paged_kv_indices_buffer=indices_buf, + paged_kv_last_page_len_buffer=last_page_buf, + ) + # plan + capture at capacity + ip, ix, lp = layout(cap_seq) + w.plan( + ip, + ix, + lp, + num_qo, + num_kv, + hd, + page, + pos_encoding_mode="NONE", + q_data_type=dtype, + kv_data_type=dtype, + ) + assert w._backend == "aiter" + for _ in range(3): + w.run(q, kv) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = w.run(q, kv) + + # replay at a shorter sequence + q.copy_(torch.randn(batch, num_qo, hd, dtype=dtype, device=device)) + ip, ix, lp = layout(replay_seq) + w.plan( + ip, + ix, + lp, + num_qo, + num_kv, + hd, + page, + pos_encoding_mode="NONE", + q_data_type=dtype, + kv_data_type=dtype, + ) + g.replay() + torch.cuda.synchronize() + + # eager AITER reference on identical inputs + ws2 = torch.zeros(64 * 1024 * 1024, dtype=torch.uint8, device=device) + ref_w = flashinfer.BatchDecodeWithPagedKVCacheWrapper(ws2, "NHD", backend="aiter") + ip, ix, lp = layout(replay_seq) + ref_w.plan( + ip, + ix, + lp, + num_qo, + num_kv, + hd, + page, + pos_encoding_mode="NONE", + q_data_type=dtype, + kv_data_type=dtype, + ) + ref = ref_w.run(q, kv) + + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)