diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index e702e4d9c..62391a0c6 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.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 diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 3491d59c8..3860df345 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.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 diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 4162e60e8..021faaef5 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.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: diff --git a/tests/buildkite/test_amd_pipeline.py b/tests/buildkite/test_amd_pipeline.py index e7c428dd1..03ea865ce 100644 --- a/tests/buildkite/test_amd_pipeline.py +++ b/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"] diff --git a/tests/buildkite/test_cuda_ar_paged_attention_lane.py b/tests/buildkite/test_cuda_ar_paged_attention_lane.py new file mode 100644 index 000000000..923743249 --- /dev/null +++ b/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" diff --git a/tests/diffusion/ar_diffusion/test_history_staging_gpu.py b/tests/diffusion/ar_diffusion/test_history_staging_gpu.py index be2776319..e04248954 100644 --- a/tests/diffusion/ar_diffusion/test_history_staging_gpu.py +++ b/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) diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 47e9b028e..05997997c 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/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")