diff --git a/tests/diffusion/ar_diffusion/test_history_staging_gpu.py b/tests/diffusion/ar_diffusion/test_history_staging_gpu.py new file mode 100644 index 000000000..be2776319 --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_history_staging_gpu.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA regressions for contiguous AR-Diffusion history staging.""" + +import pytest +import torch + +from tests.diffusion.ar_diffusion.test_paged_attention import ( + BLOCK, + HEAD_DIM, + N_HEADS, + POS, + _commit_video_span, + _gpu_flash_attn_usable, + make_state, +) +from tests.helpers.mark import hardware_test +from vllm_omni.experimental.ar_diffusion.kv_cache import paged_write_attn +from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV +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" +) +@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.""" + 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) + _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device) + stage_key, stage_value = kv.history_staging[0] + stage_key.fill_(-7) + stage_value.fill_(-7) + query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype) + for step in range(2): + ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx + ctx.max_video_tokens = 2 * BLOCK + ctx.prepare(device, action_len=0, query_len=BLOCK) + inputs = ctx.layer_inputs(0) + assert inputs.reuse_history is (step > 0) + key = torch.randn_like(query) + value = torch.randn_like(query) + actual = paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5) + expected = paged_write_attn( + inputs._replace(stage_key=None, stage_value=None, reuse_history=False), + query, + key, + value, + None, + None, + HEAD_DIM**-0.5, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + ids = inputs.block_table[0].long() + for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))): + torch.testing.assert_close( + staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0 + ) + 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.""" + 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) + _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device) + stage_key, stage_value = kv.history_staging[0] + stage_key.fill_(-7) + stage_value.fill_(-7) + query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype) + key, value = torch.randn_like(query), torch.randn_like(query) + + ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx + ctx.max_video_tokens = 2 * BLOCK + ctx.prepare(device, action_len=0, query_len=BLOCK) + inputs = ctx.layer_inputs(0) + # Populate history before testing the reuse variant, as the first denoising probe does. + paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5) + if reuse_history: + ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx + ctx.max_video_tokens = 2 * BLOCK + ctx.prepare(device, action_len=0, query_len=BLOCK) + inputs = ctx.layer_inputs(0) + assert inputs.reuse_history is reuse_history + + torch._dynamo.reset() + try: + compiled = torch.compile(paged_write_attn, mode="reduce-overhead", fullgraph=True) + for _ in range(3): + torch.compiler.cudagraph_mark_step_begin() + compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5) + current_omni_platform.synchronize() + + previous = None + for step in range(2): + key.normal_() + # Distinct value ranges ensure a stale replay cannot match the reference. + value.fill_(2.0 if step == 0 else -2.0) + torch.compiler.cudagraph_mark_step_begin() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + ) as prof: + actual = compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5) + current_omni_platform.synchronize() + assert any("cudaGraphLaunch" in event.key for event in prof.key_averages()), ( + "Inductor did not replay a CUDA graph; mutable staging inputs may have lost their static annotation" + ) + expected = paged_write_attn( + inputs._replace(stage_key=None, stage_value=None, reuse_history=False), + query, + key, + value, + None, + None, + HEAD_DIM**-0.5, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + if previous is not None: + assert not torch.equal(actual, previous) + # Copy outside capture: graph-owned output storage is reused on the next step. + previous = actual.clone() + del actual + ids = inputs.block_table[0].long() + for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))): + torch.testing.assert_close( + staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0 + ) + assert (staged[inputs.max_seq_len :] == -7).all() + finally: + torch._dynamo.reset() diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 612c974ce..7014a60c7 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/tests/diffusion/ar_diffusion/test_paged_attention.py @@ -244,7 +244,8 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, @pytest.mark.parametrize("history_chunks", [0, 1, 2, 3]) -def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks): +@pytest.mark.parametrize("window_chunks", [2, 4]) +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. The padded table always ends in at least one action-capacity block, and while the window is still @@ -254,11 +255,14 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa monkeypatch.setenv(KV_GATHER_ENV, "1") device = torch.device("cpu") dtype = torch.float32 - kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True) + kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True) + for buffer in kv.history_staging[0]: + buffer.fill_(-7) if history_chunks: _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device) ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx + ctx.max_video_tokens = 2 * BLOCK ctx.ensure_video_slots(device) key_cache, value_cache = kv.key_cache(0), kv.value_cache(0) @@ -275,6 +279,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa n_blocks = max_seq_len // BLOCK block_ids = block_table[0, :n_blocks].to(torch.long) stage_k, stage_v = ctx.history_staging(0) + # The custom op narrows the manager-owned buffers before calling this helper. + stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len] paged_attention_module._stage_window( stage_k, stage_v, @@ -287,6 +293,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa ) full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM) full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM) + for buffer in kv.history_staging[0]: + assert (buffer[max_seq_len:] == -7).all() return stage_k, stage_v, full_k, full_v, n_blocks stage_k, stage_v, full_k, full_v, n_blocks = probe(1) @@ -323,6 +331,24 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl assert (stage_k is not None) is gather_enabled +@pytest.mark.parametrize("window_chunks", [2, 4]) +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") + device = torch.device("cpu") + kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True) + stage_key, stage_value = kv.history_staging[0] + for _ in range(2): + ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx + ctx.max_video_tokens = 2 * BLOCK + ctx.prepare(device, action_len=0, query_len=BLOCK) + inputs = ctx.layer_inputs(0) + assert inputs.stage_key is stage_key + assert inputs.stage_value is stage_value + assert inputs.max_seq_len == 3 * BLOCK + assert stage_key.shape[0] == (window_chunks + 1) * BLOCK + + @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required") @pytest.mark.parametrize("history_chunks", [1, 3]) @pytest.mark.parametrize("action_len", [0, 3]) @@ -363,8 +389,8 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, query[0], current_k[0], current_v[0], - action_k[0] if action_len else None, - action_v[0] if action_len else None, + action_k[0] if action_k is not None else None, + action_v[0] if action_v is not None else None, HEAD_DIM**-0.5, ).unsqueeze(0) @@ -556,8 +582,8 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c query[0], current_k[0], current_v[0], - action_k[0] if action_len else None, - action_v[0] if action_len else None, + action_k[0] if action_k is not None else None, + action_v[0] if action_v is not None else None, HEAD_DIM**-0.5, ).unsqueeze(0) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py index 92b12f57c..2e7f72522 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py @@ -264,6 +264,7 @@ class ARDiffusionPagedForwardContext: # gather path -- their only consumer -- is switched on. if not self.kv_cache.history_staging: return + assert self.block_table is not None if action_len or self.block_table.shape[0] != 1: # Action tokens live in scratch blocks outside the video window; staging them is not modelled. return @@ -291,12 +292,14 @@ class ARDiffusionPagedForwardContext: self.stage_first_block = self.kv_len // self.block_size - self.num_current_video_blocks def history_staging(self, layer_idx: int) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """This layer's staging pair, narrowed to this forward's staged window (a prefix of the manager's buffer).""" + """Return the manager-owned tensors marked static for CUDA Graph input mutation. + + Fresh prefix views lose the bases' static-address annotation. Narrow the + window inside the attention custom op instead of at the compiled boundary. + """ if not self.staging_enabled: return None, None - key, value = self.kv_cache.history_staging[layer_idx] - tokens = int(self.max_seq_len) - return key[:tokens], value[:tokens] + return self.kv_cache.history_staging[layer_idx] def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs: if not getattr(self, "_prepared", False): @@ -435,12 +438,11 @@ def _stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, n_b earlier forward of the same AR block and cannot have changed since. """ ids = block_ids[first_block:] - blocks = n_blocks - first_block - stage_key.view(n_blocks, block_size, *key_cache.shape[2:])[first_block:].copy_( - key_cache.index_select(0, ids).view(blocks, block_size, *key_cache.shape[2:]) - ) - stage_value.view(n_blocks, block_size, *value_cache.shape[2:])[first_block:].copy_( - value_cache.index_select(0, ids).view(blocks, block_size, *value_cache.shape[2:]) + # Write directly into the caller-owned window, avoiding an intermediate + # gathered tensor and its copy on every layer and denoising step. + torch.index_select(key_cache, 0, ids, out=stage_key.view(n_blocks, block_size, *key_cache.shape[2:])[first_block:]) + torch.index_select( + value_cache, 0, ids, out=stage_value.view(n_blocks, block_size, *value_cache.shape[2:])[first_block:] ) @@ -522,7 +524,7 @@ def ar_diffusion_paged_attention( # vLLM's seqused_k/fa_version API. The ROCm flash-attn paged kernel also # requires 128-token blocks, while AR-Diffusion uses frame-aligned # 16-token blocks, so gather the visible blocks on-device first. - flash_attn_varlen_func = _rocm_flash_attn_varlen_func() + rocm_flash_attn_varlen_func = _rocm_flash_attn_varlen_func() cu_seqlens_k = torch.cat([seq_lens.new_zeros(1), torch.cumsum(seq_lens, dim=0, dtype=torch.int32)]) positions = torch.arange(int(max_seq_len), device=query_flat.device) logical_blocks = torch.div(positions, key_cache.shape[1], rounding_mode="floor") @@ -533,7 +535,7 @@ def ar_diffusion_paged_attention( valid = positions.unsqueeze(0) < seq_lens.unsqueeze(1) packed_k = gathered_k[valid] packed_v = gathered_v[valid] - out = flash_attn_varlen_func( + out = rocm_flash_attn_varlen_func( q=query_flat, k=packed_k, v=packed_v, @@ -579,8 +581,10 @@ def ar_diffusion_paged_attention( # byte-for-byte what a full gather would produce. When the history did move -- new block, # slid window, different session -- reuse_history is 0 and the whole window is re-gathered, # which is the same work the unstaged path always does. - k_flat = stage_key - v_flat = stage_value + # Keep the marked base tensors as graph inputs, but only stage the + # active prefix when the manager allocated a larger capacity. + k_flat = stage_key[:max_seq_len] + v_flat = stage_value[:max_seq_len] if reuse_history: # The caller's metadata must describe a block-aligned current chunk inside the staged # window; anything else is a wrong offset, not a reason to quietly restage everything. @@ -621,12 +625,12 @@ def ar_diffusion_paged_attention( fa_version=fa_version, ) else: - from vllm.vllm_flash_attn import flash_attn_varlen_func + from vllm.vllm_flash_attn import flash_attn_varlen_func as paged_flash_attn_varlen_func fa_version = _resolve_fa_version(query_flat.shape[-1]) out = torch.empty_like(query_flat) - flash_attn_varlen_func( + paged_flash_attn_varlen_func( q=query_flat, k=key_cache, v=value_cache,