@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 07ea9350ba

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

Signed-off-by: mobicham <hichamb@dropbox.com>

mobicham committed Sep 1, 2026 at 00:21 UTC 07ea9350baf84e33fd696d36fec9b9f24735a733

3 files changed +308 -10

tests/kernels/attention/test_triton_unified_attention.py

+253

@@ -6,8 +6,12 @@ import pytest

6 import torch

8 from vllm.platforms import current_platform

9 +from vllm.triton_utils import tl, triton

10 from vllm.utils.math_utils import next_power_of_2

11 from vllm.utils.torch_utils import set_random_seed

12 +from vllm.v1.attention.ops.triton_attention_helpers import (

13 + compute_tile_loop_bounds,

14 +)

15 from vllm.v1.attention.ops.triton_unified_attention import unified_attention

16 from vllm.v1.kv_cache_interface import KVQuantMode

@@ -32,6 +36,91 @@ NUM_BLOCKS = [32768, 2048]

36 SEQ_THRESHOLD_3D_VALUES = [0, 8]

39 +@triton.jit

40 +def _compute_clamped_mm_tile_bounds(

41 + output_ptr,

42 + mm_prefix_range_ptr,

43 + USE_CAUSAL: tl.constexpr,

44 + USE_PER_SEQ_CAUSAL: tl.constexpr,

45 + CHUNK_LOOKBACK: tl.constexpr,

46 + CHUNK_SIZE: tl.constexpr,

47 +):

48 + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(

49 + 0, # context_len

50 + 4096, # seq_len

51 + 4096, # cur_batch_query_len

52 + 140, # q_block_local_idx: query positions [1120, 1127]

53 + 0, # segm_idx_or_0

54 + 0, # tiles_per_segment_or_0

55 + 32, # TILE_SIZE

56 + 16, # BLOCK_M

57 + 8, # BLOCK_Q

58 + 2, # num_queries_per_kv

59 + 1024, # SLIDING_WINDOW

60 + True, # USE_MM_PREFIX

61 + False, # IS_3D

62 + USE_CAUSAL,

63 + USE_PER_SEQ_CAUSAL,

64 + CHUNK_LOOKBACK,

65 + CHUNK_SIZE,

66 + False, # USE_R_SWA

67 + True, # MM_PREFIX_CLAMP_SW

68 + 2, # MAX_MM_RANGES

69 + mm_prefix_range_ptr,

70 + 0, # seq_idx

71 + )

72 + tl.store(output_ptr, loop_lo)

73 + tl.store(output_ptr + 1, loop_hi)

74 + tl.store(output_ptr + 2, max_seq_prefix_len)

75 +

76 +

77 +def test_clamped_mm_prefix_prunes_sliding_window_tiles() -> None:

78 + """Retain an intersecting image range without scanning the full sequence."""

79 + mm_prefix_ranges = torch.tensor(

80 + [[[1024, 2303], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE

81 + )

82 + bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE)

83 +

84 + _compute_clamped_mm_tile_bounds[(1,)](

85 + bounds,

86 + mm_prefix_ranges,

87 + USE_CAUSAL=True,

88 + USE_PER_SEQ_CAUSAL=False,

89 + CHUNK_LOOKBACK=-1,

90 + CHUNK_SIZE=-1,

91 + )

92 +

93 + # The range is wider than the sliding window, so the upper bound must use

94 + # its inclusive endpoint rather than query_pos + sliding_window.

95 + assert bounds.tolist() == [3, 72, 4096]

96 +

97 +

98 +@pytest.mark.parametrize(

99 + ("use_causal", "use_per_seq_causal"), [(False, False), (True, True)]

100 +)

101 +def test_clamped_mm_prefix_preserves_noncausal_right_window(

102 + use_causal: bool, use_per_seq_causal: bool

103 +) -> None:

104 + """Union the non-causal right window with intersecting image ranges."""

105 + mm_prefix_ranges = torch.tensor(

106 + [[[1024, 1500], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE

107 + )

108 + bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE)

109 +

110 + _compute_clamped_mm_tile_bounds[(1,)](

111 + bounds,

112 + mm_prefix_ranges,

113 + USE_CAUSAL=use_causal,

114 + USE_PER_SEQ_CAUSAL=use_per_seq_causal,

115 + CHUNK_LOOKBACK=-1,

116 + CHUNK_SIZE=-1,

117 + )

118 +

119 + # q_hi=1127 and window=1024 admit keys through 2150, beyond the image

120 + # range endpoint at 1500. Tile 67 must therefore remain in the loop.

121 + assert bounds.tolist() == [3, 68, 4096]

122 +

123 +

124 def ref_paged_attn(

125 query: torch.Tensor,

126 key_cache: torch.Tensor,

@@ -90,6 +179,170 @@ def ref_paged_attn(

179 return torch.cat(outputs, dim=0)

180

181

182 +def ref_paged_clamped_mm_attn(

183 + query: torch.Tensor,

184 + key_cache: torch.Tensor,

185 + value_cache: torch.Tensor,

186 + query_lens: list[int],

187 + kv_lens: list[int],

188 + block_tables: torch.Tensor,

189 + mm_ranges: list[list[tuple[int, int]]],

190 + scale: float,

191 + sliding_window: int,

192 + chunk_lookback: int,

193 +) -> torch.Tensor:

194 + block_tables_cpu = block_tables.cpu().numpy()

195 + _, block_size, num_kv_heads, head_size = key_cache.shape

196 + chunk_size = sliding_window // (chunk_lookback + 1)

197 +

198 + outputs: list[torch.Tensor] = []

199 + query_start = 0

200 + for req_idx, (query_len, kv_len) in enumerate(zip(query_lens, kv_lens)):

201 + query_i = query[query_start : query_start + query_len].float()

202 + context_len = kv_len - query_len

203 + query_pos = torch.arange(query_len, device=query.device) + context_len

204 + key_pos = torch.arange(kv_len, device=query.device)

205 +

206 + num_kv_blocks = (kv_len + block_size - 1) // block_size

207 + block_indices = block_tables_cpu[req_idx, :num_kv_blocks]

208 + key_i = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]

209 + value_i = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]

210 + if query_i.shape[1] != key_i.shape[1]:

211 + repeats = query_i.shape[1] // key_i.shape[1]

212 + key_i = torch.repeat_interleave(key_i, repeats, dim=1)

213 + value_i = torch.repeat_interleave(value_i, repeats, dim=1)

214 +

215 + delta = query_pos[:, None] - key_pos[None, :]

216 + keep = (delta >= 0) & (

217 + query_pos[:, None] // chunk_size - key_pos[None, :] // chunk_size

218 + <= chunk_lookback

219 + )

220 + for range_start, range_end in mm_ranges[req_idx]:

221 + q_in_range = (query_pos >= range_start) & (query_pos <= range_end)

222 + k_in_range = (key_pos >= range_start) & (key_pos <= range_end)

223 + mm_mask = q_in_range[:, None] & k_in_range[None, :]

224 + keep |= mm_mask & (delta < sliding_window)

225 +

226 + scores = torch.einsum("qhd,khd->hqk", query_i, key_i.float()) * scale

227 + scores.masked_fill_(~keep[None], float("-inf"))

228 + probs = scores.softmax(-1).to(value_i.dtype)

229 + outputs.append(torch.einsum("hqk,khd->qhd", probs, value_i))

230 + query_start += query_len

231 +

232 + return torch.cat(outputs)

233 +

234 +

235 +@torch.inference_mode()

236 +def test_triton_unified_attn_clamped_mm_matches_dense_reference() -> None:

237 + set_random_seed(0)

238 + query_lens = [384, 96]

239 + kv_lens_list = [384, 320]

240 + mm_ranges = [[(32, 351)], [(160, 287)]]

241 + sliding_window = 128

242 + chunk_lookback = 0

243 + block_size = 16

244 + num_query_heads = 4

245 + num_kv_heads = 2

246 + head_size = 128

247 + scale = head_size**-0.5

248 +

249 + query = torch.randn(

250 + sum(query_lens),

251 + num_query_heads,

252 + head_size,

253 + dtype=torch.bfloat16,

254 + device=DEVICE_TYPE,

255 + )

256 + num_blocks_per_req = [

257 + (kv_len + block_size - 1) // block_size for kv_len in kv_lens_list

258 + ]

259 + total_blocks = sum(num_blocks_per_req)

260 + key_cache = torch.randn(

261 + total_blocks,

262 + block_size,

263 + num_kv_heads,

264 + head_size,

265 + dtype=torch.bfloat16,

266 + device=DEVICE_TYPE,

267 + )

268 + value_cache = torch.randn_like(key_cache)

269 +

270 + block_tables = torch.zeros(

271 + len(query_lens),

272 + max(num_blocks_per_req),

273 + dtype=torch.int32,

274 + device=DEVICE_TYPE,

275 + )

276 + block_start = 0

277 + for req_idx, num_blocks in enumerate(num_blocks_per_req):

278 + block_tables[req_idx, :num_blocks] = torch.arange(

279 + block_start,

280 + block_start + num_blocks,

281 + dtype=torch.int32,

282 + device=DEVICE_TYPE,

283 + )

284 + block_start += num_blocks

285 +

286 + cu_seqlens_q = torch.tensor(

287 + [0] + query_lens, dtype=torch.int32, device=DEVICE_TYPE

288 + ).cumsum(0, dtype=torch.int32)

289 + kv_lens = torch.tensor(kv_lens_list, dtype=torch.int32, device=DEVICE_TYPE)

290 + mm_prefix_range = torch.tensor(

291 + [[[32, 351]], [[160, 287]]], dtype=torch.int32, device=DEVICE_TYPE

292 + )

293 + actual = torch.empty_like(query)

294 +

295 + unified_attention(

296 + q=query,

297 + k=key_cache,

298 + v=value_cache,

299 + out=actual,

300 + cu_seqlens_q=cu_seqlens_q,

301 + max_seqlen_q=max(query_lens),

302 + seqused_k=kv_lens,

303 + max_seqlen_k=max(kv_lens_list),

304 + softmax_scale=scale,

305 + causal=True,

306 + window_size=(sliding_window - 1, 0),

307 + block_table=block_tables,

308 + softcap=0,

309 + q_descale=None,

310 + k_descale=None,

311 + v_descale=None,

312 + mm_prefix_range=mm_prefix_range,

313 + chunk_lookback=chunk_lookback,

314 + mm_prefix_clamp_sliding_window=True,

315 + )

316 +

317 + expected = ref_paged_clamped_mm_attn(

318 + query,

319 + key_cache,

320 + value_cache,

321 + query_lens,

322 + kv_lens_list,

323 + block_tables,

324 + mm_ranges,

325 + scale,

326 + sliding_window,

327 + chunk_lookback,

328 + )

329 + chunk_only = ref_paged_clamped_mm_attn(

330 + query,

331 + key_cache,

332 + value_cache,

333 + query_lens,

334 + kv_lens_list,

335 + block_tables,

336 + [[], []],

337 + scale,

338 + sliding_window,

339 + chunk_lookback,

340 + )

341 +

342 + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2)

343 + assert not torch.allclose(expected, chunk_only, atol=2e-2, rtol=2e-2)

344 +

345 +

346 @pytest.mark.parametrize(

347 "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]

348 )

vllm/v1/attention/ops/triton_attention_helpers.py

+49 -9

@@ -157,6 +157,11 @@ def compute_tile_loop_bounds(

157 USE_PER_SEQ_CAUSAL: tl.constexpr = False,

158 CHUNK_LOOKBACK: tl.constexpr = -1,

159 CHUNK_SIZE: tl.constexpr = -1,

160 + USE_R_SWA: tl.constexpr = False,

161 + MM_PREFIX_CLAMP_SW: tl.constexpr = False,

162 + MAX_MM_RANGES: tl.constexpr = 0,

163 + mm_prefix_range_ptr=None,

164 + seq_idx=0,

165 ):

166 """Compute the tile-loop bounds ``(loop_lo, loop_hi)`` and the

167 derived ``max_seq_prefix_len`` used for per-tile masking.

@@ -168,8 +173,10 @@ def compute_tile_loop_bounds(

173 mm_prefix is active or non-causal sequences need the full

174 sequence.

175 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to

171 - only tiles that can contain an allowed key under SWA.

172 - For non-causal sequences, the window extends in both directions.

176 + only tiles that can contain an allowed key under SWA. For non-causal

177 + sequences, the window extends in both directions. For Gemma4's

178 + window-clamped multimodal prefix mask, the bounds include the union of

179 + the base mask and each image range intersecting the query block.

180 3. 3D scoping: when ``IS_3D`` is True, further narrows to the

181 segment's slice via ``(segm_idx * tiles_per_segment,

182 (segm_idx + 1) * tiles_per_segment)``.

@@ -182,7 +189,7 @@ def compute_tile_loop_bounds(

189 + (BLOCK_M - 1) // num_queries_per_kv

190 + 1

191 )

185 - if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):

192 + if USE_MM_PREFIX or USE_R_SWA or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):

193 # Read the full sequence but never past seq_len: the causal-style

194 # formula above can overshoot for non-causal sequences, and slots

195 # >= seq_len are unwritten KV (last-block tail) that may hold NaN

@@ -198,8 +205,13 @@ def compute_tile_loop_bounds(

205 # Default: keep previous global behavior

206 tile_start = 0

207 tile_end = num_tiles

201 - # TODO(Isotr0py): sliding window pruning with image bidirectional mask

202 - if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:

208 + # Prefix ranges normally override the sliding window, so they require the

209 + # complete sequence. Gemma4 instead clamps prefix attention to the left

210 + # edge of the sliding window. In that case, the union of possible keys is

211 + # bounded by the normal left window and the end of any multimodal range

212 + # intersecting this query block.

213 + can_prune_sliding = (not USE_R_SWA) and ((not USE_MM_PREFIX) or MM_PREFIX_CLAMP_SW)

214 + if SLIDING_WINDOW > 0 and can_prune_sliding:

215 # Query rows covered by this Q-block

216 qpos_lo = q_block_local_idx * BLOCK_Q

217 qpos_hi = tl.minimum(

@@ -218,12 +230,40 @@ def compute_tile_loop_bounds(

230 first_allowed_key = q_abs - SLIDING_WINDOW + 1

231 if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):

232 # Non-causal: keys can be AHEAD of query within the window

221 - last_allowed_key = tl.minimum(

222 - context_len + qpos_hi + SLIDING_WINDOW - 1,

223 - seq_len - 1,

224 - )

233 + last_allowed_key = context_len + qpos_hi + SLIDING_WINDOW - 1

234 else:

235 last_allowed_key = context_len + qpos_hi

236 + if USE_MM_PREFIX and MM_PREFIX_CLAMP_SW:

237 + query_abs_lo = context_len + qpos_lo

238 + query_abs_hi = context_len + qpos_hi

239 + for i in range(MAX_MM_RANGES):

240 + range_start = tl.load(

241 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2

242 + )

243 + range_end = tl.load(

244 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1

245 + )

246 + intersects_query_block = (

247 + (range_start < range_end)

248 + & (range_start <= query_abs_hi)

249 + & (range_end >= query_abs_lo)

250 + )

251 + mm_first_allowed_key = tl.maximum(

252 + range_start, query_abs_lo - SLIDING_WINDOW + 1

253 + )

254 + first_allowed_key = tl.minimum(

255 + first_allowed_key,

256 + tl.where(

257 + intersects_query_block,

258 + mm_first_allowed_key,

259 + first_allowed_key,

260 + ),

261 + )

262 + last_allowed_key = tl.maximum(

263 + last_allowed_key,

264 + tl.where(intersects_query_block, range_end, last_allowed_key),

265 + )

266 + last_allowed_key = tl.minimum(last_allowed_key, seq_len - 1)

267 # Convert to tile indices and clamp

268 tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)

269 tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)

vllm/v1/attention/ops/triton_unified_attention.py

+6 -1

@@ -408,12 +408,17 @@ def kernel_unified_attention(

408 BLOCK_Q,

409 num_queries_per_kv,

410 SLIDING_WINDOW,

411 - USE_MM_PREFIX or USE_R_SWA,

411 + USE_MM_PREFIX,

412 IS_3D,

413 USE_CAUSAL,

414 USE_PER_SEQ_CAUSAL,

415 CHUNK_LOOKBACK,

416 CHUNK_SIZE,

417 + USE_R_SWA,

418 + MM_PREFIX_CLAMP_SW,

419 + MAX_MM_RANGES,

420 + mm_prefix_range_ptr,

421 + seq_idx,

422 )

423

424 # iterate through tiles (now limited to the sliding window range)