iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][Gemma4] Prune Triton sliding-window tiles for multimodal prefixes (#53147)

Browse Source 
Signed-off-by: mobicham <hichamb@dropbox.com>
This commit is contained in:
mobicham
2026-08-31 22:21:22 +00:00
committed by GitHub
parent d61b6e1878
commit 07ea9350ba
3 changed files with 308 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_triton_unified_attention.py
+253
View File
@@ -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)]]
)
vllm/v1/attention/ops/triton_attention_helpers.py
+49 -9
View File
@@ -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)
vllm/v1/attention/ops/triton_unified_attention.py
+6 -1
View File
@@ -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)