diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 1c1d4f803f..772fee137a 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -6,8 +6,12 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.triton_attention_helpers import ( + compute_tile_loop_bounds, +) from vllm.v1.attention.ops.triton_unified_attention import unified_attention from vllm.v1.kv_cache_interface import KVQuantMode @@ -32,6 +36,91 @@ NUM_BLOCKS = [32768, 2048] SEQ_THRESHOLD_3D_VALUES = [0, 8] +@triton.jit +def _compute_clamped_mm_tile_bounds( + output_ptr, + mm_prefix_range_ptr, + USE_CAUSAL: tl.constexpr, + USE_PER_SEQ_CAUSAL: tl.constexpr, + CHUNK_LOOKBACK: tl.constexpr, + CHUNK_SIZE: tl.constexpr, +): + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds( + 0, # context_len + 4096, # seq_len + 4096, # cur_batch_query_len + 140, # q_block_local_idx: query positions [1120, 1127] + 0, # segm_idx_or_0 + 0, # tiles_per_segment_or_0 + 32, # TILE_SIZE + 16, # BLOCK_M + 8, # BLOCK_Q + 2, # num_queries_per_kv + 1024, # SLIDING_WINDOW + True, # USE_MM_PREFIX + False, # IS_3D + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, + CHUNK_LOOKBACK, + CHUNK_SIZE, + False, # USE_R_SWA + True, # MM_PREFIX_CLAMP_SW + 2, # MAX_MM_RANGES + mm_prefix_range_ptr, + 0, # seq_idx + ) + tl.store(output_ptr, loop_lo) + tl.store(output_ptr + 1, loop_hi) + tl.store(output_ptr + 2, max_seq_prefix_len) + + +def test_clamped_mm_prefix_prunes_sliding_window_tiles() -> None: + """Retain an intersecting image range without scanning the full sequence.""" + mm_prefix_ranges = torch.tensor( + [[[1024, 2303], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE + ) + bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE) + + _compute_clamped_mm_tile_bounds[(1,)]( + bounds, + mm_prefix_ranges, + USE_CAUSAL=True, + USE_PER_SEQ_CAUSAL=False, + CHUNK_LOOKBACK=-1, + CHUNK_SIZE=-1, + ) + + # The range is wider than the sliding window, so the upper bound must use + # its inclusive endpoint rather than query_pos + sliding_window. + assert bounds.tolist() == [3, 72, 4096] + + +@pytest.mark.parametrize( + ("use_causal", "use_per_seq_causal"), [(False, False), (True, True)] +) +def test_clamped_mm_prefix_preserves_noncausal_right_window( + use_causal: bool, use_per_seq_causal: bool +) -> None: + """Union the non-causal right window with intersecting image ranges.""" + mm_prefix_ranges = torch.tensor( + [[[1024, 1500], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE + ) + bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE) + + _compute_clamped_mm_tile_bounds[(1,)]( + bounds, + mm_prefix_ranges, + USE_CAUSAL=use_causal, + USE_PER_SEQ_CAUSAL=use_per_seq_causal, + CHUNK_LOOKBACK=-1, + CHUNK_SIZE=-1, + ) + + # q_hi=1127 and window=1024 admit keys through 2150, beyond the image + # range endpoint at 1500. Tile 67 must therefore remain in the loop. + assert bounds.tolist() == [3, 68, 4096] + + def ref_paged_attn( query: torch.Tensor, key_cache: torch.Tensor, @@ -90,6 +179,170 @@ def ref_paged_attn( return torch.cat(outputs, dim=0) +def ref_paged_clamped_mm_attn( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + query_lens: list[int], + kv_lens: list[int], + block_tables: torch.Tensor, + mm_ranges: list[list[tuple[int, int]]], + scale: float, + sliding_window: int, + chunk_lookback: int, +) -> torch.Tensor: + block_tables_cpu = block_tables.cpu().numpy() + _, block_size, num_kv_heads, head_size = key_cache.shape + chunk_size = sliding_window // (chunk_lookback + 1) + + outputs: list[torch.Tensor] = [] + query_start = 0 + for req_idx, (query_len, kv_len) in enumerate(zip(query_lens, kv_lens)): + query_i = query[query_start : query_start + query_len].float() + context_len = kv_len - query_len + query_pos = torch.arange(query_len, device=query.device) + context_len + key_pos = torch.arange(kv_len, device=query.device) + + num_kv_blocks = (kv_len + block_size - 1) // block_size + block_indices = block_tables_cpu[req_idx, :num_kv_blocks] + key_i = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len] + value_i = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len] + if query_i.shape[1] != key_i.shape[1]: + repeats = query_i.shape[1] // key_i.shape[1] + key_i = torch.repeat_interleave(key_i, repeats, dim=1) + value_i = torch.repeat_interleave(value_i, repeats, dim=1) + + delta = query_pos[:, None] - key_pos[None, :] + keep = (delta >= 0) & ( + query_pos[:, None] // chunk_size - key_pos[None, :] // chunk_size + <= chunk_lookback + ) + for range_start, range_end in mm_ranges[req_idx]: + q_in_range = (query_pos >= range_start) & (query_pos <= range_end) + k_in_range = (key_pos >= range_start) & (key_pos <= range_end) + mm_mask = q_in_range[:, None] & k_in_range[None, :] + keep |= mm_mask & (delta < sliding_window) + + scores = torch.einsum("qhd,khd->hqk", query_i, key_i.float()) * scale + scores.masked_fill_(~keep[None], float("-inf")) + probs = scores.softmax(-1).to(value_i.dtype) + outputs.append(torch.einsum("hqk,khd->qhd", probs, value_i)) + query_start += query_len + + return torch.cat(outputs) + + +@torch.inference_mode() +def test_triton_unified_attn_clamped_mm_matches_dense_reference() -> None: + set_random_seed(0) + query_lens = [384, 96] + kv_lens_list = [384, 320] + mm_ranges = [[(32, 351)], [(160, 287)]] + sliding_window = 128 + chunk_lookback = 0 + block_size = 16 + num_query_heads = 4 + num_kv_heads = 2 + head_size = 128 + scale = head_size**-0.5 + + query = torch.randn( + sum(query_lens), + num_query_heads, + head_size, + dtype=torch.bfloat16, + device=DEVICE_TYPE, + ) + num_blocks_per_req = [ + (kv_len + block_size - 1) // block_size for kv_len in kv_lens_list + ] + total_blocks = sum(num_blocks_per_req) + key_cache = torch.randn( + total_blocks, + block_size, + num_kv_heads, + head_size, + dtype=torch.bfloat16, + device=DEVICE_TYPE, + ) + value_cache = torch.randn_like(key_cache) + + block_tables = torch.zeros( + len(query_lens), + max(num_blocks_per_req), + dtype=torch.int32, + device=DEVICE_TYPE, + ) + block_start = 0 + for req_idx, num_blocks in enumerate(num_blocks_per_req): + block_tables[req_idx, :num_blocks] = torch.arange( + block_start, + block_start + num_blocks, + dtype=torch.int32, + device=DEVICE_TYPE, + ) + block_start += num_blocks + + cu_seqlens_q = torch.tensor( + [0] + query_lens, dtype=torch.int32, device=DEVICE_TYPE + ).cumsum(0, dtype=torch.int32) + kv_lens = torch.tensor(kv_lens_list, dtype=torch.int32, device=DEVICE_TYPE) + mm_prefix_range = torch.tensor( + [[[32, 351]], [[160, 287]]], dtype=torch.int32, device=DEVICE_TYPE + ) + actual = torch.empty_like(query) + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=actual, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max(query_lens), + seqused_k=kv_lens, + max_seqlen_k=max(kv_lens_list), + softmax_scale=scale, + causal=True, + window_size=(sliding_window - 1, 0), + block_table=block_tables, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + mm_prefix_range=mm_prefix_range, + chunk_lookback=chunk_lookback, + mm_prefix_clamp_sliding_window=True, + ) + + expected = ref_paged_clamped_mm_attn( + query, + key_cache, + value_cache, + query_lens, + kv_lens_list, + block_tables, + mm_ranges, + scale, + sliding_window, + chunk_lookback, + ) + chunk_only = ref_paged_clamped_mm_attn( + query, + key_cache, + value_cache, + query_lens, + kv_lens_list, + block_tables, + [[], []], + scale, + sliding_window, + chunk_lookback, + ) + + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + assert not torch.allclose(expected, chunk_only, atol=2e-2, rtol=2e-2) + + @pytest.mark.parametrize( "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]] ) diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index b2667935c6..dff0b21c64 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -157,6 +157,11 @@ def compute_tile_loop_bounds( USE_PER_SEQ_CAUSAL: tl.constexpr = False, CHUNK_LOOKBACK: tl.constexpr = -1, CHUNK_SIZE: tl.constexpr = -1, + USE_R_SWA: tl.constexpr = False, + MM_PREFIX_CLAMP_SW: tl.constexpr = False, + MAX_MM_RANGES: tl.constexpr = 0, + mm_prefix_range_ptr=None, + seq_idx=0, ): """Compute the tile-loop bounds ``(loop_lo, loop_hi)`` and the derived ``max_seq_prefix_len`` used for per-tile masking. @@ -168,8 +173,10 @@ def compute_tile_loop_bounds( mm_prefix is active or non-causal sequences need the full sequence. 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to - only tiles that can contain an allowed key under SWA. - For non-causal sequences, the window extends in both directions. + only tiles that can contain an allowed key under SWA. For non-causal + sequences, the window extends in both directions. For Gemma4's + window-clamped multimodal prefix mask, the bounds include the union of + the base mask and each image range intersecting the query block. 3. 3D scoping: when ``IS_3D`` is True, further narrows to the segment's slice via ``(segm_idx * tiles_per_segment, (segm_idx + 1) * tiles_per_segment)``. @@ -182,7 +189,7 @@ def compute_tile_loop_bounds( + (BLOCK_M - 1) // num_queries_per_kv + 1 ) - if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): + if USE_MM_PREFIX or USE_R_SWA or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): # Read the full sequence but never past seq_len: the causal-style # formula above can overshoot for non-causal sequences, and slots # >= seq_len are unwritten KV (last-block tail) that may hold NaN @@ -198,8 +205,13 @@ def compute_tile_loop_bounds( # Default: keep previous global behavior tile_start = 0 tile_end = num_tiles - # TODO(Isotr0py): sliding window pruning with image bidirectional mask - if SLIDING_WINDOW > 0 and not USE_MM_PREFIX: + # Prefix ranges normally override the sliding window, so they require the + # complete sequence. Gemma4 instead clamps prefix attention to the left + # edge of the sliding window. In that case, the union of possible keys is + # bounded by the normal left window and the end of any multimodal range + # intersecting this query block. + can_prune_sliding = (not USE_R_SWA) and ((not USE_MM_PREFIX) or MM_PREFIX_CLAMP_SW) + if SLIDING_WINDOW > 0 and can_prune_sliding: # Query rows covered by this Q-block qpos_lo = q_block_local_idx * BLOCK_Q qpos_hi = tl.minimum( @@ -218,12 +230,40 @@ def compute_tile_loop_bounds( first_allowed_key = q_abs - SLIDING_WINDOW + 1 if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): # Non-causal: keys can be AHEAD of query within the window - last_allowed_key = tl.minimum( - context_len + qpos_hi + SLIDING_WINDOW - 1, - seq_len - 1, - ) + last_allowed_key = context_len + qpos_hi + SLIDING_WINDOW - 1 else: last_allowed_key = context_len + qpos_hi + if USE_MM_PREFIX and MM_PREFIX_CLAMP_SW: + query_abs_lo = context_len + qpos_lo + query_abs_hi = context_len + qpos_hi + for i in range(MAX_MM_RANGES): + range_start = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + ) + range_end = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1 + ) + intersects_query_block = ( + (range_start < range_end) + & (range_start <= query_abs_hi) + & (range_end >= query_abs_lo) + ) + mm_first_allowed_key = tl.maximum( + range_start, query_abs_lo - SLIDING_WINDOW + 1 + ) + first_allowed_key = tl.minimum( + first_allowed_key, + tl.where( + intersects_query_block, + mm_first_allowed_key, + first_allowed_key, + ), + ) + last_allowed_key = tl.maximum( + last_allowed_key, + tl.where(intersects_query_block, range_end, last_allowed_key), + ) + last_allowed_key = tl.minimum(last_allowed_key, seq_len - 1) # Convert to tile indices and clamp tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index ce4a85cfef..b0231dd0c0 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -408,12 +408,17 @@ def kernel_unified_attention( BLOCK_Q, num_queries_per_kv, SLIDING_WINDOW, - USE_MM_PREFIX or USE_R_SWA, + USE_MM_PREFIX, IS_3D, USE_CAUSAL, USE_PER_SEQ_CAUSAL, CHUNK_LOOKBACK, CHUNK_SIZE, + USE_R_SWA, + MM_PREFIX_CLAMP_SW, + MAX_MM_RANGES, + mm_prefix_range_ptr, + seq_idx, ) # iterate through tiles (now limited to the sliding window range)