@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 66b3c0e61f

[Kernel][Model] Optimize FA4 mm_prefix range lookup (#50294)

Signed-off-by: Juqi Li <2223621784@qq.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Li Juqi committed Aug 5, 2026 at 22:33 UTC 66b3c0e61f1e477820212201adf1ed871df7ee98
5 files changed +819 -45
tests/v1/attention/test_mm_prefix.py new
+636
@@ -0,0 +1,636 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""mm_prefix query-range metadata and the FA4 mask_mod that consumes it.
4 +
5 +The metadata tests run anywhere. The FA4 tests need SM100 and run the real
6 +kernel against a dense float32 reference, which is the only way to check what
7 +cannot be reasoned about from the Python side:
8 +
9 +* ``q_ranges[token_idx, 0]`` indexing an aux tensor with a runtime ``Int32``.
10 +* ``cu_seqlens_q[b] + q_local`` matching FA4's own packing of a varlen batch.
11 +
12 +The mask_mod is the *entire* mask (``(causal ∧ window) ∨ mm_prefix``), so
13 +callers must pass ``causal=False`` and no FA-layer window. FA #155 stopped
14 +auto-clearing those when ``mask_mod`` is set; leaving ``causal=True`` would
15 +short out the mask_mod on SM90 and clip bidirectional ranges on SM100.
16 +"""
17 +
18 +import numpy as np
19 +import pytest
20 +import torch
21 +
22 +from vllm.platforms import current_platform
23 +from vllm.v1.attention.backends.utils import fill_mm_prefix_query_ranges
24 +
25 +
26 +def _fa4_available() -> bool:
27 + if not current_platform.is_cuda():
28 + return False
29 + from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
30 +
31 + return is_fa_version_supported(4)
32 +
33 +
34 +requires_fa4 = pytest.mark.skipif(
35 + not _fa4_available(), reason="FA4 mm_prefix mask_mod requires FA4"
36 +)
37 +
38 +# Imported conditionally because these pull in CuTe / FA4 build artifacts that
39 +# are absent wherever requires_fa4 skips.
40 +if _fa4_available():
41 + from types import MethodType
42 +
43 + from tests.v1.attention.test_attention_backends import (
44 + MockAttentionLayer,
45 + create_and_prepopulate_kv_cache,
46 + )
47 + from tests.v1.attention.utils import (
48 + BatchSpec,
49 + create_common_attn_metadata,
50 + create_standard_kv_cache_spec,
51 + create_vllm_config,
52 + )
53 + from vllm.config import set_current_vllm_config
54 + from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
55 + from vllm.v1.attention.backends.flash_attn import (
56 + FlashAttentionImpl,
57 + FlashAttentionMetadataBuilder,
58 + _make_mm_prefix_mask_mod,
59 + )
60 + from vllm.v1.kv_cache_interface import FullAttentionSpec, get_kv_quant_mode
61 + from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
62 +
63 +DEVICE = torch.device("cuda:0")
64 +DTYPE = torch.bfloat16
65 +HEAD_SIZE = 128
66 +NUM_HEADS = 8
67 +NUM_KV_HEADS = 8
68 +SLIDING_WINDOW_LEFT = 129
69 +
70 +
71 +# --------------------------------------------------------------------------- #
72 +# Metadata construction
73 +# --------------------------------------------------------------------------- #
74 +
75 +# Stands in for the builder's persistent (max_num_batched_tokens, 2) buffer.
76 +STAGING_CAPACITY = 64
77 +
78 +
79 +def _query_ranges(mm_ranges, query_lens, seq_lens):
80 + """Fill a staging buffer and return the written rows, or None if empty."""
81 + query_start_loc = torch.tensor(
82 + [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32
83 + )
84 + # Poison the buffer so a stale row surviving the fill is visible.
85 + out = np.full((STAGING_CAPACITY, 2), 12345, dtype=np.int32)
86 + num_tokens = fill_mm_prefix_query_ranges(
87 + out,
88 + mm_ranges,
89 + query_start_loc,
90 + torch.tensor(seq_lens, dtype=torch.int32),
91 + )
92 + if num_tokens == 0:
93 + return None
94 + return torch.from_numpy(out[:num_tokens])
95 +
96 +
97 +def test_matches_range_scan_semantics_with_context_offset():
98 + """Pin the equivalence the O(1) lookup rests on.
99 +
100 + The kernel checks ``r_start <= kv_idx <= r_end`` for the range containing
101 + the query token, which is only equal to the old scan's
102 + ``any(q in r and kv in r)`` because ranges never overlap. Request 1 carries
103 + a context offset so the local-to-absolute query mapping is exercised too.
104 + """
105 + mm_ranges = {0: [(1, 3), (5, 7)], 1: [(2, 4), (9, 12)]}
106 + query_lens = [8, 6]
107 + seq_lens = [8, 13]
108 +
109 + query_ranges = _query_ranges(mm_ranges, query_lens, seq_lens)
110 + assert query_ranges is not None
111 +
112 + token_start = 0
113 + for req_idx, query_len in enumerate(query_lens):
114 + context_len = seq_lens[req_idx] - query_len
115 + for q_local in range(query_len):
116 + q_abs = context_len + q_local
117 + r_start, r_end = query_ranges[token_start + q_local].tolist()
118 + for kv_idx in range(seq_lens[req_idx]):
119 + old_scan_keep = any(
120 + start < end and start <= q_abs <= end and start <= kv_idx <= end
121 + for start, end in mm_ranges[req_idx]
122 + )
123 + new_keep = r_start <= kv_idx <= r_end
124 + assert new_keep == old_scan_keep, (req_idx, q_abs, kv_idx)
125 + token_start += query_len
126 +
127 +
128 +def test_ranges_beyond_scheduled_chunk_are_clipped():
129 + """Chunked prefill must not error or produce a wrong mask.
130 +
131 + ``disable_chunked_mm_input`` keeps a single mm item intact but still splits
132 + a prompt across steps, so a request's ranges routinely sit entirely past the
133 + tokens scheduled so far. Sizing by query token makes the out-of-chunk part
134 + a no-op rather than an out-of-bounds condition.
135 + """
136 + # Step 1 of "TTTT IIIIII": only the 4 text tokens are scheduled, so the
137 + # range is entirely ahead of the chunk. The old bounds check raised here.
138 + assert _query_ranges({0: [(4, 9)]}, query_lens=[4], seq_lens=[4]) is None
139 +
140 + # A range starting mid-chunk and running past its end keeps its absolute
141 + # bounds, so the bidirectional block still reaches the range's later keys
142 + # once they are cached. Query tokens here are absolute positions 2..5.
143 + query_ranges = _query_ranges({0: [(4, 9)]}, query_lens=[4], seq_lens=[6])
144 + assert query_ranges is not None
145 + expected = torch.tensor([[-1, -1], [-1, -1], [4, 9], [4, 9]], dtype=torch.int32)
146 + torch.testing.assert_close(query_ranges, expected)
147 +
148 +
149 +def test_buffer_reuse_does_not_leak_previous_rows():
150 + """The staging buffer persists across steps, so every reported row must be
151 + rewritten. Guards against dropping the ``-1`` fill as an optimization: a
152 + step whose ranges shrink would otherwise expose the prior step's bounds and
153 + silently widen the bidirectional mask.
154 + """
155 + out = np.zeros((STAGING_CAPACITY, 2), dtype=np.int32)
156 + query_start_loc = torch.tensor([0, 8], dtype=torch.int32)
157 + seq_lens = torch.tensor([8], dtype=torch.int32)
158 +
159 + num_tokens = fill_mm_prefix_query_ranges(
160 + out, {0: [(0, 7)]}, query_start_loc, seq_lens
161 + )
162 + assert num_tokens == 8
163 + np.testing.assert_array_equal(out[:8], np.tile([0, 7], (8, 1)))
164 +
165 + # Same rows scheduled, but now only tokens 2..3 sit inside a range.
166 + num_tokens = fill_mm_prefix_query_ranges(
167 + out, {0: [(2, 3)]}, query_start_loc, seq_lens
168 + )
169 + assert num_tokens == 8
170 + expected = np.full((8, 2), -1, dtype=np.int32)
171 + expected[2:4] = (2, 3)
172 + np.testing.assert_array_equal(out[:8], expected)
173 +
174 +
175 +def test_returns_none_when_no_range_covers_a_query_token():
176 + """No aux tensor means forward() skips the mask_mod entirely.
177 +
178 + Returning None rather than an all-``-1`` tensor is what keeps text-only and
179 + decode-only batches off the fill path.
180 + """
181 + # Text-only batch on an mm_prefix model: every request present, no ranges.
182 + assert _query_ranges({0: [], 1: []}, query_lens=[2, 2], seq_lens=[2, 2]) is None
183 + # Degenerate single-token ranges are skipped, matching the Triton path's
184 + # `start < end` validity check.
185 + assert _query_ranges({0: [(1, 1)]}, query_lens=[4], seq_lens=[4]) is None
186 + # Decode rows: the query token is generated, so it is in no range.
187 + assert _query_ranges({0: [(1, 3)]}, query_lens=[1], seq_lens=[9]) is None
188 +
189 +
190 +# --------------------------------------------------------------------------- #
191 +# FA4 kernel
192 +# --------------------------------------------------------------------------- #
193 +
194 +
195 +def _cu_seqlens(lens: list[int]) -> torch.Tensor:
196 + out = torch.zeros(len(lens) + 1, dtype=torch.int32)
197 + out[1:] = torch.tensor(lens, dtype=torch.int32).cumsum(0)
198 + return out
199 +
200 +
201 +def _dense_reference(
202 + q,
203 + k_per_req,
204 + v_per_req,
205 + query_lens,
206 + seq_lens,
207 + mm_ranges,
208 + sliding_window_left,
209 + scale,
210 + mm_clamp_sw=0,
211 +):
212 + """Dense float32 ``(causal AND window) OR mm_prefix`` per request.
213 +
214 + ``mm_clamp_sw`` confines the bidirectional block to the window, reproducing
215 + Gemma4's ``mm_prefix_clamp_sliding_window``. Only the past bound applies, so
216 + future keys inside a range still pass.
217 + """
218 + out = torch.empty_like(q, dtype=torch.float32)
219 + repeats = q.shape[1] // k_per_req[0].shape[1]
220 + q_off = 0
221 + for req_idx, (q_len, k_len) in enumerate(zip(query_lens, seq_lens)):
222 + ctx = k_len - q_len
223 + q_pos = torch.arange(q_len, device=q.device) + ctx
224 + k_pos = torch.arange(k_len, device=q.device)
225 + delta = q_pos[:, None] - k_pos[None, :]
226 +
227 + keep = delta >= 0
228 + if sliding_window_left is not None:
229 + keep &= delta < sliding_window_left
230 +
231 + for start, end in mm_ranges.get(req_idx, []):
232 + if start >= end:
233 + continue
234 + q_in = (q_pos >= start) & (q_pos <= end)
235 + k_in = (k_pos >= start) & (k_pos <= end)
236 + mm = q_in[:, None] & k_in[None, :]
237 + if mm_clamp_sw > 0:
238 + mm &= delta < mm_clamp_sw
239 + keep |= mm
240 +
241 + q_i = q[q_off : q_off + q_len].float().transpose(0, 1)
242 + k_i = k_per_req[req_idx].float().transpose(0, 1)
243 + v_i = v_per_req[req_idx].float().transpose(0, 1)
244 + if repeats > 1:
245 + k_i = k_i.repeat_interleave(repeats, dim=0)
246 + v_i = v_i.repeat_interleave(repeats, dim=0)
247 +
248 + scores = (q_i @ k_i.transpose(-1, -2)) * scale
249 + scores = scores.masked_fill(~keep[None], float("-inf"))
250 + out[q_off : q_off + q_len] = (scores.softmax(-1) @ v_i).transpose(0, 1)
251 + q_off += q_len
252 + return out
253 +
254 +
255 +def _split(packed: torch.Tensor, lens: list[int]) -> list[torch.Tensor]:
256 + return list(packed.split(lens))
257 +
258 +
259 +def _reference(q, k, v, query_lens, seq_lens, mm_ranges, sw_left, scale, **kwargs):
260 + return _dense_reference(
261 + q,
262 + _split(k, seq_lens),
263 + _split(v, seq_lens),
264 + query_lens,
265 + seq_lens,
266 + mm_ranges,
267 + sw_left,
268 + scale,
269 + **kwargs,
270 + )
271 +
272 +
273 +def _run_kernel(
274 + q, k, v, query_lens, seq_lens, mm_ranges, sliding_window_left, scale, mm_clamp_sw=0
275 +):
276 + cu_q = _cu_seqlens(query_lens)
277 + cu_k = _cu_seqlens(seq_lens)
278 +
279 + staging = torch.full((int(cu_q[-1]), 2), 12345, dtype=torch.int32).numpy()
280 + num_rows = fill_mm_prefix_query_ranges(
281 + staging, mm_ranges, cu_q, torch.tensor(seq_lens, dtype=torch.int32)
282 + )
283 + assert num_rows > 0, "test case must have at least one in-range query token"
284 + q_ranges = torch.from_numpy(staging[:num_rows]).to(DEVICE)
285 + cu_q_gpu = cu_q.to(DEVICE)
286 +
287 + out = torch.empty_like(q)
288 + flash_attn_varlen_func(
289 + q=q,
290 + k=k,
291 + v=v,
292 + out=out,
293 + cu_seqlens_q=cu_q_gpu,
294 + cu_seqlens_k=cu_k.to(DEVICE),
295 + max_seqlen_q=max(query_lens),
296 + max_seqlen_k=max(seq_lens),
297 + softmax_scale=scale,
298 + causal=False,
299 + fa_version=4,
300 + mask_mod=_make_mm_prefix_mask_mod(
301 + sliding_window=mm_clamp_sw, sliding_window_left=sliding_window_left
302 + ),
303 + aux_tensors=[q_ranges, cu_q_gpu],
304 + )
305 + return out
306 +
307 +
308 +def _randn_qkv(query_lens, seq_lens):
309 + torch.manual_seed(0)
310 + q = torch.randn(sum(query_lens), NUM_HEADS, HEAD_SIZE, dtype=DTYPE, device=DEVICE)
311 + k = torch.randn(sum(seq_lens), NUM_KV_HEADS, HEAD_SIZE, dtype=DTYPE, device=DEVICE)
312 + v = torch.randn(sum(seq_lens), NUM_KV_HEADS, HEAD_SIZE, dtype=DTYPE, device=DEVICE)
313 + return q, k, v, HEAD_SIZE**-0.5
314 +
315 +
316 +# (name, query_lens, seq_lens, mm_ranges). One shape per index computation the
317 +# mask_mod performs; shapes that differ only in request count are omitted.
318 +CASES = [
319 + # cu_seqlens_q[b] + q_local: unequal query lens across the batch, the
320 + # packing the pooling workload never produced.
321 + (
322 + "varlen_batch",
323 + [128, 64, 200],
324 + [128, 64, 200],
325 + {0: [(8, 71)], 1: [(0, 31), (40, 55)], 2: [(64, 191)]},
326 + ),
327 + # q_abs = q_idx + (seqlen_k - seqlen_q): ranges are absolute prompt
328 + # positions, so part of each range sits in the context.
329 + (
330 + "context_offset",
331 + [64],
332 + [320],
333 + {0: [(32, 95), (200, 287)]},
334 + ),
335 + # Decode rows: the query token is generated, so it is in no range and must
336 + # get the (-1, -1) sentinel. One prefill row keeps the batch mixed.
337 + (
338 + "mixed_prefill_decode",
339 + [160, 1, 1, 1],
340 + [160, 512, 300, 97],
341 + {0: [(32, 127)], 1: [(4, 259)], 2: [(8, 71)], 3: [(1, 64)]},
342 + ),
343 +]
344 +
345 +# A 320-token range, wider than the window. Every CASES range is narrower than
346 +# the window and therefore clamp-insensitive, so the window modes are exercised
347 +# on this shape instead: images larger than the window are exactly why
348 +# mm_prefix_clamp_sliding_window exists.
349 +WIDE_RANGE_CASE = ([384], [384], {0: [(32, 351)]})
350 +
351 +# (sliding_window_left, mm_clamp_sw). sliding_window_left selects between the
352 +# two mask_mod variants; mm_clamp_sw gates the Gemma4 clamp inside the sliding
353 +# one. Gemma4 sets mm_clamp_sw == sw_val, and the unclamped variant ignores
354 +# mm_clamp_sw, so a clamp without a window is not a reachable configuration.
355 +WINDOW_MODES = [
356 + (None, 0),
357 + (SLIDING_WINDOW_LEFT, 0),
358 + (SLIDING_WINDOW_LEFT, SLIDING_WINDOW_LEFT),
359 +]
360 +WINDOW_IDS = ["full_causal", "window", "window_clamped"]
361 +
362 +
363 +@requires_fa4
364 +@pytest.mark.parametrize("case", CASES, ids=[c[0] for c in CASES])
365 +def test_mm_prefix_mask_mod_matches_dense_reference(case):
366 + _, query_lens, seq_lens, mm_ranges = case
367 + q, k, v, scale = _randn_qkv(query_lens, seq_lens)
368 + args = (q, k, v, query_lens, seq_lens, mm_ranges, SLIDING_WINDOW_LEFT, scale)
369 + torch.testing.assert_close(
370 + _run_kernel(*args).float(), _reference(*args), atol=2e-2, rtol=2e-2
371 + )
372 +
373 +
374 +@requires_fa4
375 +@pytest.mark.parametrize("window_mode", WINDOW_MODES, ids=WINDOW_IDS)
376 +def test_window_modes_on_a_range_wider_than_the_window(window_mode):
377 + """Cover the three reachable window configurations.
378 +
379 + The clamped mode additionally asserts the clamp changes the mask at all,
380 + without which a range narrower than the window would leave the clamp branch
381 + unexercised while the test still passed.
382 + """
383 + sliding_window_left, mm_clamp_sw = window_mode
384 + query_lens, seq_lens, mm_ranges = WIDE_RANGE_CASE
385 + q, k, v, scale = _randn_qkv(query_lens, seq_lens)
386 + args = (q, k, v, query_lens, seq_lens, mm_ranges, sliding_window_left, scale)
387 +
388 + expected = _reference(*args, mm_clamp_sw=mm_clamp_sw)
389 + if mm_clamp_sw:
390 + assert not torch.allclose(
391 + expected, _reference(*args, mm_clamp_sw=0), atol=2e-2, rtol=2e-2
392 + ), "range is not wide enough to exercise the clamp branch"
393 +
394 + actual = _run_kernel(*args, mm_clamp_sw=mm_clamp_sw)
395 + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2)
396 +
397 +
398 +@requires_fa4
399 +def test_ranges_outside_chunk_degrade_to_causal():
400 + """A range fully behind the scheduled chunk must not widen the mask.
401 +
402 + Under chunked prefill a request's ranges routinely sit entirely in the
403 + already-computed context. Indexing by query token makes those a no-op; the
404 + result must be identical to running with no ranges at all.
405 + """
406 + query_lens, seq_lens = [64], [512]
407 + q, k, v, scale = _randn_qkv(query_lens, seq_lens)
408 +
409 + # Range (16, 79) is entirely below the chunk start (448), plus one live
410 + # range so the fill still reports rows.
411 + args = (q, k, v, query_lens, seq_lens)
412 + actual = _run_kernel(*args, {0: [(16, 79), (448, 511)]}, None, scale)
413 + expected = _reference(*args, {0: [(448, 511)]}, None, scale)
414 + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2)
415 +
416 +
417 +# --------------------------------------------------------------------------- #
418 +# FA4 through the metadata builder and the paged KV cache
419 +# --------------------------------------------------------------------------- #
420 +
421 +# Stock google/gemma-4-E2B-it ships text_config.use_bidirectional_attention as
422 +# null, so is_mm_prefix_lm is False out of the box. Tests that exercise the
423 +# builder must flip it via `_enable_mm_prefix`. Only config.json / tokenizer
424 +# are read; weights are never loaded.
425 +MODEL = "google/gemma-4-E2B-it"
426 +
427 +BLOCK_SIZE = 16
428 +SLIDING_WINDOW = 256
429 +
430 +# Request 0 is a prefill chunk whose range starts mid-context; requests 1-3 are
431 +# decodes whose generated token sits outside every range.
432 +PAGED_SEQ_LENS = [352, 513, 200, 97]
433 +PAGED_QUERY_LENS = [96, 1, 1, 1]
434 +PAGED_MM_RANGES = {
435 + 0: [(224, 287), (300, 351)],
436 + 1: [(4, 259)],
437 + 2: [(0, 63), (80, 143)],
438 + 3: [(1, 64)],
439 +}
440 +
441 +
442 +def _enable_mm_prefix(vllm_config):
443 + """Flip stock Gemma4's null use_bidirectional_attention to "vision".
444 +
445 + ``ModelConfig.__init__`` derives ``model_arch_config`` once from
446 + ``hf_text_config``, and ``is_mm_prefix_lm`` is a ``cached_property`` that
447 + may already hold False from that first derivation, so re-derive and drop
448 + the cache.
449 + """
450 + model_config = vllm_config.model_config
451 + model_config.hf_text_config.use_bidirectional_attention = "vision"
452 + model_config.model_arch_config = model_config.get_model_arch_config()
453 + model_config.__dict__.pop("is_mm_prefix_lm", None)
454 + return vllm_config
455 +
456 +
457 +@requires_fa4
458 +def test_decode_only_batch_reports_no_ranges():
459 + """Decode-only steps must carry no mm_prefix metadata at all.
460 +
461 + This is what makes FULL CUDA graph capture consistent for this feature.
462 + Capture runs through ``_dummy_run`` with no multimodal requests, so no
463 + mask_mod is attached to the captured graph; a decode-only replay must
464 + therefore also want no mask_mod, or the captured graph would silently drop
465 + it. Indexing by query token gives that for free, because a generated token
466 + is always past every prompt range.
467 +
468 + The range-id scheme keyed off the request having ranges at all, so it
469 + attached a mask_mod to decode steps and did not have this property.
470 + """
471 + vllm_config = _enable_mm_prefix(
472 + create_vllm_config(
473 + model_name=MODEL,
474 + max_model_len=1024,
475 + block_size=BLOCK_SIZE,
476 + num_gpu_blocks=256,
477 + )
478 + )
479 + decode_batch = BatchSpec(
480 + seq_lens=[513, 200, 97], query_lens=[1, 1, 1], name="decode_only"
481 + )
482 + common = create_common_attn_metadata(decode_batch, BLOCK_SIZE, DEVICE)
483 + common.mm_req_doc_ranges = {0: [(4, 259)], 1: [(0, 63)], 2: [(1, 64)]}
484 +
485 + builder = FlashAttentionMetadataBuilder(
486 + create_standard_kv_cache_spec(vllm_config),
487 + ["model.layers.0.self_attn.attn"],
488 + vllm_config,
489 + DEVICE,
490 + )
491 + md = builder.build(common_prefix_len=0, common_attn_metadata=common)
492 + assert md.mm_prefix_query_range_tensor is None
493 +
494 +
495 +@requires_fa4
496 +@pytest.mark.parametrize(
497 + "head_size",
498 + [
499 + # 128: FA4's default SM100 path (not the dedicated hd256 kernel), so
500 + # paged KV + seqused_k works on B200 and H100. Non-paged mm_prefix
501 + # cases in this file already use HEAD_SIZE=128.
502 + pytest.param(128, id="hd128"),
503 + # 512: Gemma4 global head dim; FA4 resolves on H100, skips on
504 + # Blackwell (TMEM).
505 + pytest.param(512, id="global_512"),
506 + ],
507 +)
508 +def test_mm_prefix_kv_cache_path(head_size: int):
509 + """FA4 + mask_mod + block_table matches a dense reference.
510 +
511 + Use head_size=128 on B200/H100 (FA4 + paged/`seqused_k` supported) and
512 + 512 on H100 only (Blackwell FA4 rejects 512 via TMEM). Avoid 256: SM100's
513 + dedicated hd256 2CTA kernel asserts away ``seqused_k``, which paged KV
514 + requires.
515 +
516 + Tensors, kv_cache_spec, and FlashAttentionImpl must all use the same
517 + ``head_size`` — ``run_attention_backend`` / ``create_standard_kv_cache_spec``
518 + read ``get_head_size()`` (=512 for Gemma4), so this test builds them
519 + explicitly. The sliding window is set on the impl; mm_prefix encodes it in
520 + mask_mod and clears FA's built-in window at forward time.
521 + """
522 + torch.manual_seed(0)
523 + vllm_config = _enable_mm_prefix(
524 + create_vllm_config(
525 + model_name=MODEL,
526 + max_model_len=max(PAGED_SEQ_LENS),
527 + block_size=BLOCK_SIZE,
528 + num_gpu_blocks=2048,
529 + )
530 + )
531 + assert vllm_config.model_config.is_mm_prefix_lm
532 +
533 + mc = vllm_config.model_config
534 + # Keep every get_head_size() reader (builder helpers, etc.) on the same dim.
535 + mc.get_head_size = MethodType(lambda self: head_size, mc)
536 +
537 + # Skip under the same config Gemma4 forces (flash_attn_version=4), so the
538 + # check matches FlashAttentionImpl's get_flash_attn_version call.
539 + with set_current_vllm_config(vllm_config):
540 + if get_flash_attn_version(head_size=head_size) != 4:
541 + pytest.skip(f"FA4 does not support head_size={head_size} on this device")
542 +
543 + num_heads = mc.get_num_attention_heads(vllm_config.parallel_config)
544 + num_kv_heads = mc.get_num_kv_heads(vllm_config.parallel_config)
545 + scale = head_size**-0.5
546 +
547 + batch = BatchSpec(
548 + seq_lens=PAGED_SEQ_LENS, query_lens=PAGED_QUERY_LENS, name="mm_prefix_mixed"
549 + )
550 + kv_cache_spec = FullAttentionSpec(
551 + block_size=BLOCK_SIZE,
552 + num_kv_heads=num_kv_heads,
553 + head_size=head_size,
554 + dtype=mc.dtype,
555 + sliding_window=SLIDING_WINDOW,
556 + kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
557 + )
558 + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
559 + common.mm_req_doc_ranges = PAGED_MM_RANGES
560 +
561 + qs, new_ks, new_vs, k_fulls, v_fulls, k_ctxs, v_ctxs = [], [], [], [], [], [], []
562 + for q_len, s_len in zip(PAGED_QUERY_LENS, PAGED_SEQ_LENS):
563 + ctx = s_len - q_len
564 + shape = (num_kv_heads, head_size)
565 + qs.append(torch.randn(q_len, num_heads, head_size, dtype=DTYPE, device=DEVICE))
566 + k_full = torch.randn(s_len, *shape, dtype=DTYPE, device=DEVICE)
567 + v_full = torch.randn(s_len, *shape, dtype=DTYPE, device=DEVICE)
568 + k_fulls.append(k_full)
569 + v_fulls.append(v_full)
570 + k_ctxs.append(k_full[:ctx])
571 + v_ctxs.append(v_full[:ctx])
572 + new_ks.append(k_full[ctx:])
573 + new_vs.append(v_full[ctx:])
574 +
575 + query = torch.cat(qs)
576 + key = torch.cat(new_ks)
577 + value = torch.cat(new_vs)
578 + kv_cache = create_and_prepopulate_kv_cache(
579 + k_contexts=k_ctxs,
580 + v_contexts=v_ctxs,
581 + block_size=BLOCK_SIZE,
582 + num_kv_heads=num_kv_heads,
583 + head_size=head_size,
584 + dtype=DTYPE,
585 + device=DEVICE,
586 + num_blocks=2048,
587 + common_attn_metadata=common,
588 + randomize_blocks=True,
589 + )
590 +
591 + # Packed KV layout is [..., 2 * head_size]; impl.split(head_size) needs match.
592 + assert kv_cache.shape[-1] == 2 * head_size
593 +
594 + layer_names = ["model.layers.0.self_attn.attn"]
595 + with set_current_vllm_config(vllm_config):
596 + builder = FlashAttentionMetadataBuilder(
597 + kv_cache_spec, layer_names, vllm_config, DEVICE
598 + )
599 + attn_metadata = builder.build(common_prefix_len=0, common_attn_metadata=common)
600 + # head_size=128 keeps FA4 under local attention on Blackwell (the
601 + # FA4→FA2 demotion is only for local+256). Window is still encoded in
602 + # mask_mod at forward; FA's built-in window is cleared there.
603 + impl = FlashAttentionImpl(
604 + num_heads=num_heads,
605 + head_size=head_size,
606 + scale=scale,
607 + num_kv_heads=num_kv_heads,
608 + alibi_slopes=None,
609 + sliding_window=SLIDING_WINDOW,
610 + kv_cache_dtype="auto",
611 + )
612 + assert impl.vllm_flash_attn_version == 4, (
613 + f"expected FA4, got FA{impl.vllm_flash_attn_version} for "
614 + f"head_size={head_size}"
615 + )
616 + assert impl.head_size == head_size == kv_cache_spec.head_size
617 +
618 + mock_layer = MockAttentionLayer(DEVICE)
619 + output = torch.empty_like(query)
620 + impl.do_kv_cache_update(
621 + mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
622 + )
623 + actual = impl.forward(
624 + mock_layer, query, key, value, kv_cache, attn_metadata, output=output
625 + )
626 +
627 + ref_args = (query, k_fulls, v_fulls, PAGED_QUERY_LENS, PAGED_SEQ_LENS)
628 + expected = _dense_reference(*ref_args, PAGED_MM_RANGES, SLIDING_WINDOW, scale)
629 + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2)
630 +
631 + # Without this the assertion above would also pass if the mask_mod were
632 + # silently dropped, since the batch would then be plain causal + window.
633 + causal_only = _dense_reference(*ref_args, {}, SLIDING_WINDOW, scale)
634 + assert not torch.allclose(expected, causal_only, atol=2e-2, rtol=2e-2), (
635 + "test batch does not actually exercise the mm_prefix branch"
636 + )
vllm/model_executor/models/gemma4_mm.py
+2
@@ -2078,6 +2078,8 @@ class Gemma4ForConditionalGeneration(
2078 metadata.mm_prefix_range = None
2079 if hasattr(metadata, "mm_prefix_range_tensor"):
2080 metadata.mm_prefix_range_tensor = None
2081 + if hasattr(metadata, "mm_prefix_query_range_tensor"):
2082 + metadata.mm_prefix_query_range_tensor = None
2083
2084 if isinstance(attn_metadata, list):
2085 for ub_metadata in attn_metadata:
vllm/v1/attention/backend.py
+1 -1
@@ -471,7 +471,7 @@ class CommonAttentionMetadata:
471 """PrefixLM bidirectional ranges for multimodal tokens. Maps
472 request index to list of (start, end) token position ranges
473 where bidirectional attention should apply. None for text-only
474 - batches or non-PrefixLM models."""
474 + batches or non-PrefixLM models. A request's ranges must not overlap."""
475
476 rswa_prefix_lens: torch.Tensor | None = None
477 """(batch_size,) per-request prefix length (prompt/image token count) for
vllm/v1/attention/backends/flash_attn.py
+109 -44
@@ -3,6 +3,7 @@
3 """Attention layer with FlashAttention."""
4
5 import copy
6 +import functools
7 from dataclasses import dataclass
8 from typing import ClassVar
9
@@ -12,6 +13,7 @@ import torch
13 from vllm.model_executor.layers.attention import Attention
14 from vllm.platforms import current_platform
15 from vllm.utils.torch_utils import (
16 + PIN_MEMORY,
17 canonicalize_singleton_dim_strides,
18 is_quantized_kv_cache,
19 )
@@ -28,7 +30,10 @@ from vllm.v1.attention.backends.fa_utils import (
30 is_fa_version_supported,
31 is_flash_attn_varlen_func_available,
32 )
31 -from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
33 +from vllm.v1.attention.backends.utils import (
34 + fill_mm_prefix_query_ranges,
35 + get_dcp_local_seq_lens,
36 +)
37 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
38 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
39 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
@@ -284,9 +289,10 @@ class FlashAttentionMetadata:
289
290 sliding_window: tuple[int, int] | None = None
291
287 - # PrefixLM bidirectional ranges for multimodal tokens.
288 - # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.
289 - mm_prefix_range_tensor: torch.Tensor | None = None
292 + # PrefixLM bidirectional range containing each scheduled query token.
293 + # Shape: (num_actual_tokens, 2) int32, absolute [start, end] bounds;
294 + # (-1, -1) for query tokens outside every multimodal range.
295 + mm_prefix_query_range_tensor: torch.Tensor | None = None
296
297 # Reference Sliding Window Attention (R-SWA) fields.
298 # rswa_prefix_lens: per-request prompt lengths [num_reqs], int32, CUDA.
@@ -455,6 +461,21 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
461 [self.rswa_window], dtype=torch.int32, device=self.device
462 )
463
464 + # mm_prefix: persistent staging + device buffers owned by this builder,
465 + # sized by scheduled query tokens so build() never allocates.
466 + self.mm_prefix_query_ranges_cpu: torch.Tensor | None = None
467 + self.mm_prefix_query_ranges_np: np.ndarray | None = None
468 + self.mm_prefix_query_ranges_gpu: torch.Tensor | None = None
469 + if self.model_config.is_mm_prefix_lm:
470 + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
471 + self.mm_prefix_query_ranges_cpu = torch.empty(
472 + (max_num_tokens, 2), dtype=torch.int32, pin_memory=PIN_MEMORY
473 + )
474 + self.mm_prefix_query_ranges_np = self.mm_prefix_query_ranges_cpu.numpy()
475 + self.mm_prefix_query_ranges_gpu = torch.empty(
476 + (max_num_tokens, 2), dtype=torch.int32, device=self.device
477 + )
478 +
479 def build(
480 self,
481 common_prefix_len: int,
@@ -696,16 +717,31 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
717 )
718
719 # Compute mm_prefix range tensor if the batch contains
699 - # multimodal tokens with bidirectional ranges.
720 + # multimodal tokens with bidirectional ranges. Built for every FA
721 + # group; Gemma4 nulls the field for its non-sliding layers.
722 mm_ranges = common_attn_metadata.mm_req_doc_ranges
701 - if mm_ranges is not None:
702 - from vllm.v1.attention.backends.utils import (
703 - compute_mm_prefix_range_tensor,
723 + if mm_ranges is not None and self.mm_prefix_query_ranges_np is not None:
724 + # The upper bound is exact for prefill rows, which is where
725 + # mm_prefix ranges live; decode rows only ever get an optimistic
726 + # (larger) context, moving them further past every range.
727 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None, (
728 + "mm_prefix requires seq_lens_cpu_upper_bound"
729 )
705 -
706 - attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor(
707 - mm_ranges, num_reqs, seq_lens.device
730 + num_mm_tokens = fill_mm_prefix_query_ranges(
731 + self.mm_prefix_query_ranges_np,
732 + mm_ranges,
733 + common_attn_metadata.query_start_loc_cpu,
734 + common_attn_metadata.seq_lens_cpu_upper_bound,
735 )
736 + if num_mm_tokens > 0:
737 + assert self.mm_prefix_query_ranges_cpu is not None
738 + assert self.mm_prefix_query_ranges_gpu is not None
739 + mm_query_ranges = self.mm_prefix_query_ranges_gpu[:num_mm_tokens]
740 + mm_query_ranges.copy_(
741 + self.mm_prefix_query_ranges_cpu[:num_mm_tokens],
742 + non_blocking=True,
743 + )
744 + attn_metadata.mm_prefix_query_range_tensor = mm_query_ranges
745
746 # R-SWA: copy prefix lengths into persistent buffers (outside the
747 # compiled region) so forward() never allocates during CUDA graph
@@ -971,23 +1007,25 @@ class FlashAttentionImpl(AttentionImpl):
1007 causal = attn_metadata.causal
1008 is_dynamic_causal = isinstance(causal, torch.Tensor)
1009
974 - mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor
1010 + mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
1011 mm_mask_mod = None
1012 mm_aux = None
1013 if (
978 - mm_prefix_ranges is not None
1014 + mm_prefix_query_ranges is not None
1015 and not is_dynamic_causal
1016 and causal is True
1017 and self.vllm_flash_attn_version == 4
1018 ):
983 - max_ranges = mm_prefix_ranges.shape[1]
984 - # Sliding window value in Triton convention
985 - # (1 + window_size[0]). Global-attention layers
986 - # store (-1, -1) → sw stays None / 0.
1019 + # Use the layer impl's window, not attn_metadata's. The
1020 + # metadata field comes from kv_cache_spec (model-wide, e.g.
1021 + # Gemma4's 512), while the impl holds the per-layer window
1022 + # the mask_mod must encode (e.g. a test override).
1023 + # Triton convention: 1 + window_size[0]. Global layers store
1024 + # (-1, -1) → sw stays None.
1025 + layer_window = self.sliding_window
1026 sw_val = (
988 - 1 + sliding_window_size[0]
989 - if sliding_window_size is not None
990 - and sliding_window_size[0] >= 0
1027 + 1 + layer_window[0]
1028 + if layer_window is not None and layer_window[0] >= 0
1029 else None
1030 )
1031 # Gemma4: also clamp the bidirectional block to the
@@ -1000,11 +1038,16 @@ class FlashAttentionImpl(AttentionImpl):
1038 ):
1039 mm_clamp_sw = sw_val
1040 mm_mask_mod = _make_mm_prefix_mask_mod(
1003 - max_ranges,
1041 sliding_window=mm_clamp_sw,
1042 sliding_window_left=sw_val,
1043 )
1007 - mm_aux = [mm_prefix_ranges]
1044 + mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc]
1045 + # mm_prefix is (causal ∧ window) ∨ bidirectional-range —
1046 + # not ⊆ causal. FA #155 stopped auto-clearing causal/local
1047 + # when mask_mod is set, so the caller must disable them or
1048 + # the built-in causal path shorts out / clips the mask_mod.
1049 + causal = False
1050 + sliding_window_size = None
1051
1052 # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
1053 # mask without block-size approximation. The mask_mod encodes
@@ -1378,14 +1421,19 @@ class FlashAttentionImpl(AttentionImpl):
1421 return output
1422
1423
1424 +@functools.cache
1425 def _make_mm_prefix_mask_mod(
1382 - max_ranges: int,
1426 sliding_window: int = 0,
1427 sliding_window_left: int | None = None,
1428 ):
1429 """Build a CuTE-DSL mask_mod implementing
1430 ``(causal AND sliding_window) OR mm_prefix``.
1431
1432 + Cached so identical ``(sliding_window, sliding_window_left)`` reuse the
1433 + same function object. FA4's ``hash_callable`` mixes ``repr()`` of closure
1434 + cells into the compile key; the nested ``_load_q_range`` would otherwise
1435 + get a new address each call and force a full JIT recompile every forward.
1436 +
1437 The FA4 kernel passes *local* ``q_idx`` (0-based within the current
1438 prefill chunk) while ``kv_idx`` is absolute (0-based over the full
1439 KV cache). We recover the absolute Q position via
@@ -1394,6 +1442,15 @@ def _make_mm_prefix_mask_mod(
1442 use consistent absolute positions. This matches the Triton
1443 reference path (``compute_kv_seq_mask``).
1444
1445 + ``aux_tensors[0]`` holds the absolute ``[start, end]`` bounds of the
1446 + mm_prefix range containing each scheduled query token (``(-1, -1)`` when
1447 + none), and ``aux_tensors[1]`` is ``cu_seqlens_q``, used to turn the local
1448 + ``q_idx`` into a packed row index. Because mm_prefix ranges never overlap,
1449 + ``r_start <= kv_idx <= r_end`` is exactly "query and key share a range", so
1450 + no key-side lookup is needed and ``kv_idx`` is never used as an index.
1451 + The ``(-1, -1)`` sentinel falls out for free: ``kv_idx <= -1`` is false for
1452 + every valid key.
1453 +
1454 ``sliding_window_left`` enforces the sliding window on the causal
1455 term (None = full causal, no window). ``sliding_window`` clamps the
1456 bidirectional block to the window (0 = unclamped; >0 = Gemma4 local
@@ -1405,8 +1462,29 @@ def _make_mm_prefix_mask_mod(
1462
1463 from vllm.vllm_flash_attn.cute.utils import ( # type: ignore[import-untyped]
1464 scalar_to_ssa,
1465 + ssa_to_scalar,
1466 )
1467
1468 + @cute.jit
1469 + def _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx):
1470 + """Load the mm_prefix range bounds for this query row.
1471 +
1472 + Both loads depend only on the query index, so they hoist out of the
1473 + unrolled per-element mask loop. ``ssa_to_scalar`` reads lane 0, so one
1474 + call must not span query rows (hence the ``__vec_size__`` pin below).
1475 + Clamping keeps ``token_idx`` in bounds for partial tiles and padded
1476 + (``seqlen_q == 0``) rows; neither produces output.
1477 + """
1478 + q_ranges = aux_tensors[0]
1479 + cu_seqlens_q = aux_tensors[1]
1480 + b = batch_idx[0]
1481 + q_local = cutlass.min(ssa_to_scalar(q_idx), seqlen_info.seqlen_q - Int32(1))
1482 + token_idx = cutlass.max(cu_seqlens_q[b] + q_local, Int32(0))
1483 + return (
1484 + scalar_to_ssa(q_ranges[token_idx, 0], Int32),
1485 + scalar_to_ssa(q_ranges[token_idx, 1], Int32),
1486 + )
1487 +
1488 if sliding_window_left is not None:
1489
1490 @cute.jit
@@ -1422,18 +1500,11 @@ def _make_mm_prefix_mask_mod(
1500 q_abs = q_idx + ctx_off
1501 sw = scalar_to_ssa(Int32(sliding_window_left), Int32)
1502 keep = (kv_idx <= q_abs) & ((q_abs - kv_idx) < sw)
1425 - ranges = aux_tensors[0]
1426 - b = batch_idx[0]
1427 - for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1428 - r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1429 - r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1430 - valid = r_start < r_end
1431 - q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid
1432 - k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1433 - mm = q_in & k_in
1434 - if sliding_window > 0:
1435 - mm = mm & ((q_abs - kv_idx) < sw)
1436 - keep = keep | mm
1503 + r_start, r_end = _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx)
1504 + mm = (kv_idx >= r_start) & (kv_idx <= r_end)
1505 + if sliding_window > 0:
1506 + mm = mm & ((q_abs - kv_idx) < sw)
1507 + keep = keep | mm
1508 return keep
1509
1510 else:
@@ -1450,18 +1521,12 @@ def _make_mm_prefix_mask_mod(
1521 ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
1522 q_abs = q_idx + ctx_off
1523 keep = kv_idx <= q_abs
1453 - ranges = aux_tensors[0]
1454 - b = batch_idx[0]
1455 - for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1456 - r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1457 - r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1458 - valid = r_start < r_end
1459 - q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid
1460 - k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1461 - keep = keep | (q_in & k_in)
1524 + r_start, r_end = _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx)
1525 + keep = keep | ((kv_idx >= r_start) & (kv_idx <= r_end))
1526 return keep
1527
1528 mm_prefix_mask_mod.use_fast_sampling = True
1529 + mm_prefix_mask_mod.__vec_size__ = 1 # _load_q_range takes lane 0 of q_idx
1530 return mm_prefix_mask_mod
1531
1532
vllm/v1/attention/backends/utils.py
+71
@@ -75,6 +75,77 @@ def compute_mm_prefix_range_tensor(
75 return padded.view(num_seqs, max_ranges, 2)
76
77
78 +def fill_mm_prefix_query_ranges(
79 + out: np.ndarray,
80 + mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
81 + query_start_loc_cpu: torch.Tensor,
82 + seq_lens_cpu: torch.Tensor,
83 +) -> int:
84 + """Map each scheduled query token to the mm_prefix range containing it.
85 +
86 + Writes into ``out``, a caller-owned ``(max_num_batched_tokens, 2)`` int32
87 + staging buffer, and returns the number of rows written (0 if no range
88 + covers any scheduled query token, in which case ``out`` is untouched and
89 + the caller should skip the mask_mod entirely). Row ``i`` holds the absolute
90 + ``[start, end]`` bounds of the bidirectional range that query token ``i``
91 + belongs to, or ``(-1, -1)`` when it is outside every range.
92 +
93 + mm_prefix ranges never overlap, so "query and key share a range" is
94 + equivalent to "the key lies inside the query's own range". The kernel
95 + therefore needs no key-side lookup, and this metadata is sized by scheduled
96 + query tokens rather than by context length -- bounded by
97 + ``max_num_batched_tokens`` instead of ``num_seqs * max_seq_len``.
98 +
99 + Ranges are absolute prompt positions and may extend past the tokens
100 + scheduled so far under chunked prefill; the portion outside the current
101 + chunk is simply not recorded. Degenerate ranges (``start >= end``) are
102 + skipped to match the Triton path's ``start < end`` validity check.
103 +
104 + ``seq_lens_cpu`` only needs to be exact for prefill rows, since mm_prefix
105 + ranges cover prompt tokens: an over-estimate on a decode row shifts that
106 + row's query position further past every range, which still matches nothing.
107 + """
108 + if mm_prefix_range is None:
109 + return 0
110 +
111 + query_start_loc = query_start_loc_cpu.numpy()
112 + num_actual_tokens = int(query_start_loc[-1])
113 + if num_actual_tokens <= 0:
114 + return 0
115 + assert num_actual_tokens <= out.shape[0], (
116 + f"mm_prefix staging buffer holds {out.shape[0]} tokens, got {num_actual_tokens}"
117 + )
118 +
119 + # Resolve every span before touching `out`, so batches whose ranges all
120 + # fall outside the scheduled tokens skip the fill entirely.
121 + spans: list[tuple[int, int, int, int]] = []
122 + for req_idx, req_ranges in mm_prefix_range.items():
123 + if not req_ranges:
124 + continue
125 + token_start = int(query_start_loc[req_idx])
126 + query_len = int(query_start_loc[req_idx + 1]) - token_start
127 + if query_len <= 0:
128 + continue
129 + # Absolute position of this request's first scheduled query token.
130 + context_len = int(seq_lens_cpu[req_idx]) - query_len
131 + for start, end in req_ranges:
132 + if start >= end:
133 + continue
134 + first = max(start - context_len, 0)
135 + last = min(end - context_len, query_len - 1)
136 + if first > last:
137 + continue
138 + spans.append((token_start + first, token_start + last + 1, start, end))
139 +
140 + if not spans:
141 + return 0
142 +
143 + out[:num_actual_tokens] = -1
144 + for row_start, row_end, start, end in spans:
145 + out[row_start:row_end] = (start, end)
146 + return num_actual_tokens
147 +
148 +
149 def is_valid_kv_cache_layout(value: str) -> bool:
150 return value in get_args(KVCacheLayoutType)
151