ci: route AR paged attention tests to GPU lanes (#8346)

Signed-off-by: andyluo7 <andy.luo@amd.com>

.buildkite/amd/test-amd-merge.yml

@@ -91,6 +91,25 @@ steps:
 
 - group: ":card_index_dividers: Diffusion Test"
   steps:
+    # AR paged attention is real GPU work, including cold AITER compilation.
+    # Keep its complete numerical matrix out of the CPU-marked Simple shards.
+    - label: "ROCm · AR Diffusion Paged Attention GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 20
+      commands:
+        - export VLLM_ROCM_USE_AITER=1
+        - export VLLM_OMNI_AR_FA_REQUIRED=1
+        - export AITER_JIT_DIR="/tmp/vllm-omni-ar-aiter-$$BUILDKITE_JOB_ID"
+        - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-ar-torch-extensions-$$BUILDKITE_JOB_ID"
+        - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+        - >-
+          timeout --signal=TERM --kill-after=2m 15m
+          pytest -sv tests/diffusion/ar_diffusion/test_paged_attention.py
+          -m 'core_model and rocm and MI325 and cards_1' --run-level core_model
+
     - label: "Diffusion · LoRA Test"
       agent_pool: mi300_1
       depends_on: amd-build

.buildkite/amd/test-amd-ready.yml

@@ -129,6 +129,25 @@ steps:
           -m 'core_model and rocm and MI325 and cards_1'
           --run-level "core_model"
 
+    # AR paged attention is real GPU work, including cold AITER compilation.
+    # Keep its complete numerical matrix out of the CPU-marked Simple shards.
+    - label: "ROCm · AR Diffusion Paged Attention GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 20
+      commands:
+        - export VLLM_ROCM_USE_AITER=1
+        - export VLLM_OMNI_AR_FA_REQUIRED=1
+        - export AITER_JIT_DIR="/tmp/vllm-omni-ar-aiter-$$BUILDKITE_JOB_ID"
+        - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-ar-torch-extensions-$$BUILDKITE_JOB_ID"
+        - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+        - >-
+          timeout --signal=TERM --kill-after=2m 15m
+          pytest -sv tests/diffusion/ar_diffusion/test_paged_attention.py
+          -m 'core_model and rocm and MI325 and cards_1' --run-level core_model
+
     - label: "Diffusion · Batch Test"
       agent_pool: mi300_1
       depends_on: amd-build

.buildkite/cuda/test-merge.yml

@@ -38,6 +38,17 @@ steps:
   - group: ":card_index_dividers: Diffusion Test"
     depends_on: upload-merge-pipeline
     steps:
+      - label: "Diffusion · AR Paged Attention GPU Test"
+        timeout_in_minutes: 20
+        commands:
+          - export VLLM_OMNI_AR_FA_REQUIRED=1
+          - >-
+            timeout 15m pytest -sv
+            tests/diffusion/ar_diffusion/test_paged_attention.py
+            -m 'core_model and cuda and L4 and cards_1'
+            --run-level "core_model"
+        mirror_hardwares: l4_1
+
       - label: "Diffusion · LoRA Test"
         timeout_in_minutes: 30
         commands:

tests/buildkite/test_amd_pipeline.py

@@ -13,6 +13,9 @@ AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
 AMD_NIGHTLY_PIPELINE = Path(".buildkite/amd/test-amd-nightly.yml")
 AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
 AMD_TEMPLATE = Path(".buildkite/amd/test-template-amd-omni.j2")
+AR_PAGED_ATTENTION_LABEL = "ROCm · AR Diffusion Paged Attention GPU Test"
+DIFFUSION_GROUP = ":card_index_dividers: Diffusion Test"
+AR_PAGED_ATTENTION_MARKERS = "core_model and rocm and MI325 and cards_1"
 
 
 def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict:
@@ -31,6 +34,35 @@ def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict:
     return step
 
 
+def _find_group(label: str, pipeline_path: Path) -> dict:
+    pipeline = yaml.safe_load(pipeline_path.read_text(encoding="utf-8"))
+    group = next((step for step in pipeline.get("steps", []) if step.get("group") == label), None)
+    assert group is not None, f"missing AMD pipeline group: {label}"
+    return group
+
+
+def test_ar_paged_attention_gpu_lane_is_blocking_and_pinned() -> None:
+    lane_definitions = []
+    for pipeline_path in (AMD_READY_PIPELINE, AMD_MERGE_PIPELINE):
+        group = _find_group(DIFFUSION_GROUP, pipeline_path)
+        step = next((item for item in group["steps"] if item.get("label") == AR_PAGED_ATTENTION_LABEL), None)
+        assert step is not None, f"missing {AR_PAGED_ATTENTION_LABEL} from {DIFFUSION_GROUP} in {pipeline_path}"
+        lane_definitions.append(step)
+
+        assert step["grade"] == "Blocking"
+        assert step["timeout_in_minutes"] == 20
+        assert "export VLLM_OMNI_AR_FA_REQUIRED=1" in step["commands"]
+
+        pytest_command = next(command for command in step["commands"] if "pytest" in command)
+        argv = split(pytest_command)
+        assert argv[:4] == ["timeout", "--signal=TERM", "--kill-after=2m", "15m"]
+        assert "tests/diffusion/ar_diffusion/test_paged_attention.py" in argv
+        assert argv[argv.index("-m") + 1] == AR_PAGED_ATTENTION_MARKERS
+        assert argv[argv.index("--run-level") + 1] == "core_model"
+
+    assert lane_definitions[0] == lane_definitions[1]
+
+
 def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None:
     step = _find_step("Qwen3-TTS Base E2E Test")
     commands = step["commands"]

tests/buildkite/test_cuda_ar_paged_attention_lane.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+from shlex import split
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+CUDA_MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
+DIFFUSION_GROUP = ":card_index_dividers: Diffusion Test"
+AR_PAGED_ATTENTION_LABEL = "Diffusion · AR Paged Attention GPU Test"
+
+
+def test_cuda_merge_keeps_ar_paged_attention_gpu_coverage() -> None:
+    pipeline = yaml.safe_load(CUDA_MERGE_PIPELINE.read_text(encoding="utf-8"))
+    group = next(step for step in pipeline["steps"] if step.get("group") == DIFFUSION_GROUP)
+    lane = next(step for step in group["steps"] if step.get("label") == AR_PAGED_ATTENTION_LABEL)
+
+    assert lane["timeout_in_minutes"] == 20
+    assert lane["mirror_hardwares"] == "l4_1"
+    assert "export VLLM_OMNI_AR_FA_REQUIRED=1" in lane["commands"]
+
+    pytest_command = next(command for command in lane["commands"] if "pytest" in command)
+    argv = split(pytest_command)
+    assert argv[:2] == ["timeout", "15m"]
+    assert "tests/diffusion/ar_diffusion/test_paged_attention.py" in argv
+    assert argv[argv.index("-m") + 1] == "core_model and cuda and L4 and cards_1"
+    assert argv[argv.index("--run-level") + 1] == "core_model"

tests/diffusion/ar_diffusion/test_history_staging_gpu.py

@@ -11,7 +11,7 @@ from tests.diffusion.ar_diffusion.test_paged_attention import (
     N_HEADS,
     POS,
     _commit_video_span,
-    _gpu_flash_attn_usable,
+    _require_gpu_flash_attn,
     make_state,
 )
 from tests.helpers.mark import hardware_test
@@ -22,14 +22,18 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
-@pytest.mark.skipif(
-    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
-)
+def _require_cuda_flash_attn() -> None:
+    if torch.version.hip is not None:
+        pytest.skip("usable CUDA FlashAttention is required")
+    _require_gpu_flash_attn()
+
+
 @hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
 @pytest.mark.parametrize("window_chunks", [2, 4])
 @torch.inference_mode()
 def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window_chunks):
     """Full and reused gathers match fresh attention without touching spare capacity or stale history."""
+    _require_cuda_flash_attn()
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device, dtype = torch.device("cuda"), torch.bfloat16
     kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
@@ -65,15 +69,13 @@ def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window
             assert (staged[inputs.max_seq_len :] == -7).all()
 
 
-@pytest.mark.skipif(
-    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
-)
 @hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
 @pytest.mark.parametrize("window_chunks", [2, 4])
 @pytest.mark.parametrize("reuse_history", [False, True])
 @torch.inference_mode()
 def test_staged_attention_replays_inductor_cudagraph(monkeypatch, window_chunks, reuse_history):
     """Inductor must capture the real mutable op and replay with fresh current K/V."""
+    _require_cuda_flash_attn()
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device, dtype = torch.device("cuda"), torch.bfloat16
     kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import os
 import subprocess
 from importlib.util import find_spec
 from pathlib import Path
@@ -12,6 +13,7 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
@@ -26,7 +28,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
 BLOCK = 16
@@ -118,6 +120,17 @@ def _gpu_flash_attn_usable() -> bool:
     return driver_major >= 580
 
 
+def _require_gpu_flash_attn() -> None:
+    """Skip optional local runs, but never let a dedicated GPU lane pass vacuously."""
+    if _gpu_flash_attn_usable():
+        return
+
+    reason = "usable GPU FlashAttention is required"
+    if os.environ.get("VLLM_OMNI_AR_FA_REQUIRED"):
+        pytest.fail(f"{reason} when VLLM_OMNI_AR_FA_REQUIRED is set")
+    pytest.skip(reason)
+
+
 def _commit_video_span(
     kv: ARDiffusionKVCache,
     st: ARDiffusionKVState,
@@ -139,6 +152,7 @@ def _commit_video_span(
     return k, v
 
 
+@pytest.mark.cpu
 def test_paged_context_allocates_lazily_and_commits_after_forward():
     _, st = make_state()
 
@@ -157,6 +171,7 @@ def test_paged_context_allocates_lazily_and_commits_after_forward():
     assert st._committed[POS] == BLOCK
 
 
+@pytest.mark.cpu
 def test_scratch_video_and_action_blocks_do_not_commit():
     kv, st = make_state()
 
@@ -171,6 +186,7 @@ def test_scratch_video_and_action_blocks_do_not_commit():
     assert st._committed[POS] == 0
 
 
+@pytest.mark.cpu
 def test_pipeline_kv_get_paged_path_has_no_gather_backend():
     kv, st = make_state()
     assert not hasattr(kv, "gather_window_all_layers")
@@ -188,6 +204,7 @@ def test_pipeline_kv_get_paged_path_has_no_gather_backend():
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 @pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.cpu
 def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
     torch.manual_seed(0)
     device = torch.device("cpu")
@@ -264,6 +281,7 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
 
 @pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
 @pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.cpu
 def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
     """A second probe of the same AR block restages its current K/V where the table actually holds it.
 
@@ -330,6 +348,7 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
 
 
 @pytest.mark.parametrize("gather_enabled", [False, True])
+@pytest.mark.cpu
 def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
     if gather_enabled:
         monkeypatch.setenv(KV_GATHER_ENV, "1")
@@ -351,6 +370,7 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl
 
 
 @pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.cpu
 def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
     """Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
     monkeypatch.setenv(KV_GATHER_ENV, "1")
@@ -368,11 +388,12 @@ def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks
         assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("history_chunks", [1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 @pytest.mark.parametrize("commit_current", [False, True])
 def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    _require_gpu_flash_attn()
     torch.manual_seed(0)
     device = torch.device("cuda")
     dtype = torch.float16
@@ -439,6 +460,7 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
     torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
 
 
+@pytest.mark.cpu
 def test_block_table_padded_to_fixed_width():
     """Shapes must be constant across window growth: only values change."""
     device = torch.device("cpu")
@@ -462,6 +484,7 @@ def test_block_table_padded_to_fixed_width():
     assert int(ctx2.seq_lens[0]) == 2 * BLOCK
 
 
+@pytest.mark.cpu
 def test_prepare_is_idempotent_and_layers_share_metadata():
     device = torch.device("cpu")
     kv, st = make_state(num_layers=2)
@@ -487,6 +510,7 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
     assert i1.value_pool is kv._v_pools[1]
 
 
+@pytest.mark.cpu
 def test_layer_inputs_before_prepare_raises():
     _, st = make_state()
     layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
@@ -494,6 +518,7 @@ def test_layer_inputs_before_prepare_raises():
         layer_ctx.to_layer_inputs()
 
 
+@pytest.mark.cpu
 def test_custom_op_registration_idempotent():
     import importlib
     import sys
@@ -508,6 +533,7 @@ def test_custom_op_registration_idempotent():
     assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
 
 
+@pytest.mark.cpu
 def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
     # Older PyTorch ADInplaceOrView handlers index positional mutable inputs
     # directly. Default-valued trailing inputs can disappear before that handler.
@@ -518,6 +544,7 @@ def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
 
 
 @pytest.mark.parametrize("reuse_history_staging", [False, True])
+@pytest.mark.cpu
 def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
     """The op must trace as one opaque node: fullgraph OK, and changed tensor
     VALUES (new slots / block ids) must not trigger recompilation."""
@@ -576,6 +603,7 @@ RAGGED_CHUNK = 24
 
 @pytest.mark.parametrize("commit_current", [False, True])
 @pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.cpu
 def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
     """Attention reads a sequence, not a set of blocks.
 
@@ -643,6 +671,7 @@ def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action
     torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
 
 
+@pytest.mark.cpu
 def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
     """The write targets themselves must leave no hole.
 
@@ -668,6 +697,7 @@ def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
     assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids
 
 
+@pytest.mark.cpu
 def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
     """The table's shape must not track how full the window is.
 
@@ -769,6 +799,7 @@ def _model_block_indices(kv, adapter) -> dict[int, int]:
     ("sink_chunks", "window_chunks", "reset_at_boundary"),
     [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
 )
+@pytest.mark.cpu
 def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
     sink_chunks, window_chunks, reset_at_boundary
 ):
@@ -859,6 +890,7 @@ def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
     ("sink_chunks", "window_chunks", "reset_at_boundary"),
     [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
 )
+@pytest.mark.cpu
 def test_compaction_does_not_change_which_tokens_stay_resident(
     monkeypatch, sink_chunks, window_chunks, reset_at_boundary
 ):
@@ -898,6 +930,7 @@ def test_compaction_does_not_change_which_tokens_stay_resident(
         assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"
 
 
+@pytest.mark.cpu
 def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
     """The staging buffers must fit the table width this module pads to.
 
@@ -925,6 +958,7 @@ def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch
 
 
 @pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.cpu
 def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
     """Action K/V follows the video blocks in the table, and the kernel reads it as one run.
 
@@ -939,6 +973,7 @@ def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_cur
         ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))
 
 
+@pytest.mark.cpu
 def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
     """832x480 with the checkpoint's own sink and window, slot by slot.
 
@@ -988,7 +1023,7 @@ def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
         kv_len = int(seq_lens[0])
 
         position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
-        read = []
+        read: list[tuple[int, int]] = []
         for b in table[0].tolist()[: -(-kv_len // block)]:
             assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
             read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
@@ -1039,6 +1074,7 @@ class _CountedBlock:
         return self._block.block_id
 
 
+@pytest.mark.cpu
 def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
     """However long the block table is, reading the window must not cost more.
 
@@ -1074,6 +1110,7 @@ def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
     )
 
 
+@pytest.mark.cpu
 def test_the_checkpoints_own_default_resolution_can_build_a_cache():
     """832x480 is what LingBot World v2 ships as its default, and it could not run.
 
@@ -1111,6 +1148,7 @@ def test_the_checkpoints_own_default_resolution_can_build_a_cache():
     assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98
 
 
+@pytest.mark.cpu
 def test_the_shipped_resolution_geometry_matches_dense_attention():
     """The real number, not a scaled-down stand-in: 1560 tokens per frame.
 
@@ -1163,7 +1201,7 @@ def test_the_shipped_resolution_geometry_matches_dense_attention():
     torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("commit_current", [False, True])
 def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     """The shipped 832x480 geometry through FlashAttention's paged kernel.
@@ -1173,6 +1211,7 @@ def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     tokens per frame, through the production path -- host prep followed by
     the fused write+attend op, which is what a forward actually calls.
     """
+    _require_gpu_flash_attn()
     torch.manual_seed(0)
     device = torch.device("cuda")
     dtype = torch.float16
@@ -1204,7 +1243,7 @@ def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
@@ -1213,6 +1252,7 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c
     Covers an empty, partial and full window plus a partially filled action block:
     the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
     """
+    _require_gpu_flash_attn()
 
     torch.manual_seed(0)
     device = torch.device("cuda")