@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 74a1ac38b0
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 07:30 UTC

[v1] Add PrefixLM support to TritonAttention backend (#30386)

Isotr0py committed Dec 18, 2025 at 08:05 UTC 74a1ac38b00a8cf502db085d1bbd77712cf47e41
4 files changed +281 -124
tests/models/multimodal/generation/test_multimodal_gguf.py
+99 -34
@@ -1,17 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Literal, NamedTuple
4 +import os
5 +
6 +os.environ["TOKENIZERS_PARALLELISM"] = "true"
7 +
8 +from typing import Any, NamedTuple
9
10 import pytest
11 from huggingface_hub import hf_hub_download
12 from pytest import MarkDecorator
13 +from transformers import AutoModelForImageTextToText
14
15 from tests.quantization.utils import is_quant_method_supported
16 from vllm.assets.image import ImageAsset
17 +from vllm.multimodal.image import rescale_image_size
18 from vllm.utils.torch_utils import set_default_torch_num_threads
19
14 -from ....conftest import PromptImageInput, VllmRunner
20 +from ....conftest import IMAGE_ASSETS, HfRunner, VllmRunner
21 from ...utils import check_logprobs_close
22
23
@@ -21,9 +27,10 @@ class GGUFMMTestConfig(NamedTuple):
27 gguf_backbone: str
28 gguf_mmproj: str
29 prompt: list[str]
24 - mm_data: dict[Literal["images"], PromptImageInput]
30 + image_names: list[str] # Store names, load PIL images at runtime
31 max_model_len: int = 4096
32 marks: list[MarkDecorator] = []
33 + mm_processor_kwargs: dict[str, Any] = {}
34
35 @property
36 def gguf_model(self):
@@ -31,27 +38,75 @@ class GGUFMMTestConfig(NamedTuple):
38 return hf_hub_download(self.gguf_repo, filename=self.gguf_backbone)
39
40
41 +# Common prompts aligned with test_common.py "gemma3" entry format
42 +_GEMMA3_PROMPTS = IMAGE_ASSETS.prompts(
43 + {
44 + "stop_sign": (
45 + "<bos><start_of_turn>user\n"
46 + "<start_of_image>What's the content in the center of the image?"
47 + "<end_of_turn>\n<start_of_turn>model\n"
48 + ),
49 + "cherry_blossom": (
50 + "<bos><start_of_turn>user\n"
51 + "<start_of_image>What is the season?"
52 + "<end_of_turn>\n<start_of_turn>model\n"
53 + ),
54 + }
55 +)
56 +
57 +# Image asset names - load at runtime to avoid pickle issues with subprocess
58 +_GEMMA3_IMAGE_NAMES = ["stop_sign", "cherry_blossom"]
59 +
60 +# Regular multimodal (no pan-and-scan) - uses QAT Q4_0 GGUF
61 GEMMA3_CONFIG = GGUFMMTestConfig(
62 original_model="google/gemma-3-4b-it",
63 gguf_repo="google/gemma-3-4b-it-qat-q4_0-gguf",
64 gguf_backbone="gemma-3-4b-it-q4_0.gguf",
65 gguf_mmproj="mmproj-model-f16-4B.gguf",
39 - prompt=["<start_of_image>Describe this image in detail:"],
40 - mm_data={"images": [ImageAsset("stop_sign").pil_image]},
66 + prompt=_GEMMA3_PROMPTS,
67 + image_names=_GEMMA3_IMAGE_NAMES,
68 + max_model_len=4096,
69 + marks=[pytest.mark.core_model],
70 + mm_processor_kwargs={},
71 +)
72 +
73 +# Pan-and-scan multimodal - uses unquantized BF16 GGUF
74 +GEMMA3_CONFIG_PAN_AND_SCAN = GGUFMMTestConfig(
75 + original_model="google/gemma-3-4b-it",
76 + gguf_repo="unsloth/gemma-3-4b-it-GGUF",
77 + gguf_backbone="gemma-3-4b-it-BF16.gguf",
78 + gguf_mmproj="mmproj-BF16.gguf",
79 + prompt=_GEMMA3_PROMPTS,
80 + image_names=_GEMMA3_IMAGE_NAMES,
81 + max_model_len=4096,
82 marks=[pytest.mark.core_model],
83 + mm_processor_kwargs={"do_pan_and_scan": True},
84 )
85
44 -MODELS_TO_TEST = [GEMMA3_CONFIG]
86 +MODELS_TO_TEST = [GEMMA3_CONFIG, GEMMA3_CONFIG_PAN_AND_SCAN]
87
88
89 def run_multimodal_gguf_test(
90 + hf_runner: type[HfRunner],
91 vllm_runner: type[VllmRunner],
92 model: GGUFMMTestConfig,
93 dtype: str,
94 max_tokens: int,
95 num_logprobs: int,
96 ):
54 - # Run gguf model.
97 + # Load images at runtime (inside subprocess) to avoid pickle issues
98 + images = [ImageAsset(name).pil_image for name in model.image_names]
99 + size_factors = [0.25, 0.5, 1.0]
100 + inputs_per_image = [
101 + (
102 + [prompt for _ in size_factors],
103 + [rescale_image_size(image, factor) for factor in size_factors],
104 + )
105 + for image, prompt in zip(images, model.prompt)
106 + ]
107 +
108 + # NOTE: Run vLLM first to avoid CUDA init issues with multiprocessing fork.
109 + # Run GGUF model via vLLM.
110 with (
111 set_default_torch_num_threads(1),
112 vllm_runner(
@@ -60,36 +115,43 @@ def run_multimodal_gguf_test(
115 tokenizer_name=model.original_model,
116 dtype=dtype,
117 max_model_len=model.max_model_len,
118 + mm_processor_kwargs=model.mm_processor_kwargs,
119 ) as gguf_model,
120 ):
65 - gguf_outputs = gguf_model.generate_greedy_logprobs(
66 - prompts=model.prompt,
67 - max_tokens=max_tokens,
68 - num_logprobs=num_logprobs,
69 - **model.mm_data,
70 - )
71 -
72 - # Run unquantized model.
73 - with vllm_runner(
74 - model_name=model.original_model,
75 - enforce_eager=True, # faster tests
121 + gguf_outputs_per_case = [
122 + gguf_model.generate_greedy_logprobs(
123 + prompts,
124 + max_tokens,
125 + num_logprobs=num_logprobs,
126 + images=images,
127 + )
128 + for prompts, images in inputs_per_image
129 + ]
130 +
131 + # Then run HfRunner for HuggingFace baseline comparison.
132 + with hf_runner(
133 + model.original_model,
134 dtype=dtype,
77 - max_model_len=model.max_model_len,
78 - ) as original_model:
79 - original_outputs = original_model.generate_greedy_logprobs(
80 - prompts=model.prompt,
81 - max_tokens=max_tokens,
82 - num_logprobs=num_logprobs,
83 - **model.mm_data,
135 + auto_cls=AutoModelForImageTextToText,
136 + ) as hf_model:
137 + hf_outputs_per_case = [
138 + hf_model.generate_greedy_logprobs_limit(
139 + prompts,
140 + max_tokens,
141 + num_logprobs=num_logprobs,
142 + images=images,
143 + )
144 + for prompts, images in inputs_per_image
145 + ]
146 +
147 + for hf_outputs, gguf_outputs in zip(hf_outputs_per_case, gguf_outputs_per_case):
148 + check_logprobs_close(
149 + outputs_0_lst=hf_outputs,
150 + outputs_1_lst=gguf_outputs,
151 + name_0="hf",
152 + name_1="gguf",
153 )
154
86 - check_logprobs_close(
87 - outputs_0_lst=original_outputs,
88 - outputs_1_lst=gguf_outputs,
89 - name_0="original",
90 - name_1="gguf",
91 - )
92 -
155
156 @pytest.mark.skipif(
157 not is_quant_method_supported("gguf"),
@@ -105,11 +167,14 @@ def run_multimodal_gguf_test(
167 @pytest.mark.parametrize("dtype", ["bfloat16"])
168 @pytest.mark.parametrize("max_tokens", [32])
169 @pytest.mark.parametrize("num_logprobs", [10])
108 -def test_models(
170 +def test_gemma3_mm_gguf(
171 + hf_runner: type[HfRunner],
172 vllm_runner: type[VllmRunner],
173 model: GGUFMMTestConfig,
174 dtype: str,
175 max_tokens: int,
176 num_logprobs: int,
177 ) -> None:
115 - run_multimodal_gguf_test(vllm_runner, model, dtype, max_tokens, num_logprobs)
178 + run_multimodal_gguf_test(
179 + hf_runner, vllm_runner, model, dtype, max_tokens, num_logprobs
180 + )
vllm/attention/ops/triton_unified_attention.py
+143 -21
@@ -86,6 +86,9 @@ def kernel_unified_attention_2d(
86 USE_SOFTCAP: tl.constexpr, # bool
87 USE_SINKS: tl.constexpr, # bool
88 SLIDING_WINDOW: tl.constexpr, # int
89 + USE_MM_PREFIX: tl.constexpr, # bool
90 + MAX_MM_RANGES: tl.constexpr, # int
91 + mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence
92 stride_k_cache_0: tl.int64, # int
93 stride_k_cache_1: tl.int64, # int
94 stride_k_cache_2: tl.int64, # int
@@ -270,7 +273,38 @@ def kernel_unified_attention_2d(
273 else:
274 V = V_load
275
273 - seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1
276 + # Compute attention mask: causal by default (key <= query)
277 + query_abs_pos = context_len + query_pos[:, None]
278 + seq_mask = seq_offset[None, :] <= query_abs_pos
279 +
280 + # Apply sliding window to base mask BEFORE mm_prefix OR.
281 + # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix
282 + if SLIDING_WINDOW > 0:
283 + seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
284 +
285 + # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
286 + # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
287 + if USE_MM_PREFIX:
288 + for i in range(MAX_MM_RANGES):
289 + range_start = tl.load(
290 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2
291 + )
292 + range_end = tl.load(
293 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1
294 + )
295 +
296 + is_valid = range_start < range_end
297 + q_in_range = (
298 + (query_abs_pos >= range_start)
299 + & (query_abs_pos <= range_end)
300 + & is_valid
301 + )
302 + k_in_range = (
303 + (seq_offset[None, :] >= range_start)
304 + & (seq_offset[None, :] <= range_end)
305 + & is_valid
306 + )
307 + seq_mask |= q_in_range & k_in_range
308
309 # S : (BLOCK_M, TILE_SIZE)
310 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
@@ -284,13 +318,6 @@ def kernel_unified_attention_2d(
318 query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
319 )
320
287 - if SLIDING_WINDOW > 0:
288 - S = tl.where(
289 - (context_len + query_pos[:, None] - seq_offset) < SLIDING_WINDOW,
290 - S,
291 - float("-inf"),
292 - )
293 -
321 if USE_ALIBI_SLOPES:
322 S += alibi_slope[:, None] * (seq_offset - context_len)
323
@@ -398,6 +425,9 @@ def kernel_unified_attention_3d(
425 num_seqs: tl.int32,
426 BLOCK_M: tl.constexpr, # int
427 NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
428 + USE_MM_PREFIX: tl.constexpr, # bool
429 + MAX_MM_RANGES: tl.constexpr, # int
430 + mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence
431 ):
432 q_block_global_idx = tl.program_id(0)
433 kv_head_idx = tl.program_id(1)
@@ -559,7 +589,38 @@ def kernel_unified_attention_3d(
589 else:
590 V = V_load
591
562 - seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1
592 + # Compute attention mask: causal by default (key <= query)
593 + query_abs_pos = context_len + query_pos[:, None]
594 + seq_mask = seq_offset[None, :] <= query_abs_pos
595 +
596 + # Apply sliding window to base mask BEFORE mm_prefix OR.
597 + # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix
598 + if SLIDING_WINDOW > 0:
599 + seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
600 +
601 + # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
602 + # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
603 + if USE_MM_PREFIX:
604 + for i in range(MAX_MM_RANGES):
605 + range_start = tl.load(
606 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2
607 + )
608 + range_end = tl.load(
609 + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1
610 + )
611 +
612 + is_valid = range_start < range_end
613 + q_in_range = (
614 + (query_abs_pos >= range_start)
615 + & (query_abs_pos <= range_end)
616 + & is_valid
617 + )
618 + k_in_range = (
619 + (seq_offset[None, :] >= range_start)
620 + & (seq_offset[None, :] <= range_end)
621 + & is_valid
622 + )
623 + seq_mask |= q_in_range & k_in_range
624
625 # S : (BLOCK_M, TILE_SIZE)
626 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
@@ -572,13 +633,6 @@ def kernel_unified_attention_3d(
633 query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
634 )
635
575 - if SLIDING_WINDOW > 0:
576 - S = tl.where(
577 - (context_len + query_pos[:, None] - seq_offset) < SLIDING_WINDOW,
578 - S,
579 - float("-inf"),
580 - )
581 -
636 if USE_ALIBI_SLOPES:
637 S += alibi_slope[:, None] * (seq_offset - context_len)
638
@@ -732,6 +786,43 @@ def reduce_segments(
786 tl.store(output_ptr + output_offset, acc, mask=dim_mask)
787
788
789 +def _is_gemma3_attention(head_size: int, sliding_window: int) -> bool:
790 + """Detect Gemma3 models via unique (head_size, sliding_window) signature.
791 +
792 + Gemma3 models are the only ones using sliding_window=1024 with
793 + head_size 128 (27B) or 256 (1B, 4B, 12B). Other SWA models use
794 + different window sizes (Mistral=4096, Phi-3=2047).
795 + """
796 + return sliding_window == 1024 and head_size in (128, 256)
797 +
798 +
799 +def _get_tile_size(
800 + head_size: int,
801 + sliding_window: int,
802 + element_size: int,
803 + is_mm_prefix: bool,
804 + is_prefill: bool,
805 +) -> int:
806 + """Select tile size with Gemma3-specific optimization.
807 +
808 + For Gemma3, use 32 for both prefill and decode to better utilize
809 + the larger head dimension (128/256). For other models, use
810 + the default vLLM behavior.
811 + """
812 + if is_mm_prefix:
813 + # Multimodal bidirectional attention needs a larger tile size
814 + return 64
815 +
816 + if _is_gemma3_attention(head_size, sliding_window):
817 + # Gemma3: use 32 for decode (default is 16)
818 + return 32
819 +
820 + # Default behavior
821 + if is_prefill:
822 + return 32
823 + return 16 if element_size >= 2 else 32
824 +
825 +
826 def unified_attention(
827 q,
828 k,
@@ -759,6 +850,8 @@ def unified_attention(
850 qq_bias=None,
851 # Optional tensor for sinks
852 sinks=None,
853 + # Optional tensor for prefix lengths (PrefixLM support)
854 + mm_prefix_range=None,
855 ):
856 assert causal, "Only causal attention is supported"
857 assert q_descale is None, "Q scales not supported"
@@ -766,6 +859,17 @@ def unified_attention(
859 if sinks is not None:
860 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
861
862 + use_mm_prefix = False
863 + max_mm_ranges = 0
864 + if mm_prefix_range is not None:
865 + if mm_prefix_range.ndim == 3:
866 + use_mm_prefix = True
867 + max_mm_ranges = mm_prefix_range.shape[1]
868 + else:
869 + raise ValueError(
870 + f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
871 + )
872 +
873 use_alibi_slopes = alibi_slopes is not None
874 use_qq_bias = qq_bias is not None
875
@@ -792,11 +896,23 @@ def unified_attention(
896 # = floor(q.shape[0] / BLOCK_Q) + num_seqs
897 total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
898
795 - # Assigning default tile sizes for prefill and decode.
796 - # Note: each tile size must be at least 32 for "fp8" (q.element_size() == 1)
797 - # and at least 16 for all other data types.
798 - TILE_SIZE_PREFILL = 32
799 - TILE_SIZE_DECODE = 16 if q.element_size() >= 2 else 32
899 + # Tile sizes for prefill and decode. Gemma3 models use optimized values.
900 + # Note: tile size must be at least 32 for fp8 (element_size == 1).
901 + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
902 + TILE_SIZE_PREFILL = _get_tile_size(
903 + head_size,
904 + sliding_window_val,
905 + q.element_size(),
906 + is_mm_prefix=use_mm_prefix,
907 + is_prefill=True,
908 + )
909 + TILE_SIZE_DECODE = _get_tile_size(
910 + head_size,
911 + sliding_window_val,
912 + q.element_size(),
913 + is_mm_prefix=use_mm_prefix,
914 + is_prefill=False,
915 + )
916
917 # Launch the 2D kernel if
918 # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
@@ -847,6 +963,9 @@ def unified_attention(
963 USE_QQ_BIAS=use_qq_bias,
964 USE_SOFTCAP=(softcap > 0),
965 USE_SINKS=(sinks is not None),
966 + USE_MM_PREFIX=use_mm_prefix,
967 + MAX_MM_RANGES=max_mm_ranges,
968 + mm_prefix_range_ptr=mm_prefix_range,
969 SLIDING_WINDOW=(1 + window_size[0]),
970 stride_k_cache_0=k.stride(0),
971 stride_k_cache_1=k.stride(1),
@@ -895,6 +1014,9 @@ def unified_attention(
1014 USE_QQ_BIAS=use_qq_bias,
1015 USE_SOFTCAP=(softcap > 0),
1016 USE_SINKS=(sinks is not None),
1017 + USE_MM_PREFIX=use_mm_prefix,
1018 + MAX_MM_RANGES=max_mm_ranges,
1019 + mm_prefix_range_ptr=mm_prefix_range,
1020 SLIDING_WINDOW=(1 + window_size[0]),
1021 stride_k_cache_0=k.stride(0),
1022 stride_k_cache_1=k.stride(1),
vllm/model_executor/models/gemma3.py
-69
@@ -19,7 +19,6 @@ from collections.abc import Iterable
19 from itertools import islice
20
21 import torch
22 -import torch.nn.functional as F
22 from torch import nn
23 from transformers import Gemma3TextConfig
24
@@ -226,77 +225,9 @@ class Gemma3Attention(nn.Module):
225
226 q, k = self.rotary_emb(positions, q, k)
227 attn_output = self.attn(q, k, v)
229 -
230 - if not kwargs.get("has_images", False):
231 - # Fast path for text-only inputs. The performance for the text-only
232 - # inputs are not affected by the naive attention below.
233 - output, _ = self.o_proj(attn_output)
234 - return output
235 -
236 - # NOTE(woosuk): Gemma3 uses bidirectional attention between image tokens
237 - # that correspond to the same image while using causal attention
238 - # otherwise. Current attention backends cannot handle this pattern, so
239 - # we temporarily use a naive attention implementation with mask tensors.
240 -
241 - # We intentionally keep the attention backend as-is and only override
242 - # `attn_output` with the naive implementation's output. This minimizes
243 - # changes to existing model runners and attention backends. The call to
244 - # `self.attn(q, k, v)` is only used to populate the KV cache - its
245 - # output is discarded and overwritten below. While this duplicates
246 - # computation, it maintains compatibility.
247 - # TODO(woosuk): Optimize by implementing custom attention kernels.
248 - attn_output = self.naive_attn_with_masks(q, k, v, out=attn_output, **kwargs)
228 output, _ = self.o_proj(attn_output)
229 return output
230
252 - def naive_attn_with_masks(
253 - self,
254 - q: torch.Tensor,
255 - k: torch.Tensor,
256 - v: torch.Tensor,
257 - out: torch.Tensor,
258 - **kwargs,
259 - ) -> torch.Tensor:
260 - # NOTE(woosuk): As described in the comment above, this code is not
261 - # meant to be performant. It is only meant to be correct.
262 - q = q.view(-1, self.num_heads, self.head_dim)
263 - # Expand the key and value to handle GQA.
264 - num_queries_per_kv = self.num_heads // self.num_kv_heads
265 - k = k.view(-1, self.num_kv_heads, self.head_dim)
266 - k = k.repeat_interleave(num_queries_per_kv, dim=-2)
267 - v = v.view(-1, self.num_kv_heads, self.head_dim)
268 - v = v.repeat_interleave(num_queries_per_kv, dim=-2)
269 -
270 - if self.is_sliding:
271 - attn_masks = kwargs["local_attn_masks"]
272 - else:
273 - attn_masks = kwargs["global_attn_masks"]
274 -
275 - seq_lens = kwargs["seq_lens"]
276 - start_idx = 0
277 - for seq_len, attn_mask in zip(seq_lens, attn_masks):
278 - end_idx = start_idx + seq_len
279 - query = q[start_idx:end_idx].unsqueeze(0)
280 - key = k[start_idx:end_idx].unsqueeze(0)
281 - value = v[start_idx:end_idx].unsqueeze(0)
282 -
283 - # Transpose.
284 - query = query.transpose(1, 2)
285 - key = key.transpose(1, 2)
286 - value = value.transpose(1, 2)
287 -
288 - output = F.scaled_dot_product_attention(
289 - query,
290 - key,
291 - value,
292 - attn_mask,
293 - self.scaling,
294 - )
295 - output = output.transpose(1, 2).flatten(-2, -1)
296 - out[start_idx:end_idx] = output
297 - start_idx = end_idx
298 - return out
299 -
231
232 class Gemma3DecoderLayer(nn.Module):
233 def __init__(
vllm/v1/attention/backends/triton_attn.py
+39
@@ -76,6 +76,39 @@ class TritonAttentionMetadata:
76 # Optional aot scheduling
77 scheduler_metadata: torch.Tensor | None = None
78 prefix_scheduler_metadata: torch.Tensor | None = None
79 + mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
80 +
81 + @property
82 + def mm_prefix_range_tensor(self) -> torch.Tensor | None:
83 + """Convert mm_prefix_range dict to padded tensor for Triton kernel.
84 +
85 + Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges.
86 + Empty ranges have start==end==0, which kernel skips via is_valid check.
87 + """
88 + # TODO(Isotr0py): Move to model runner's attention metadata
89 + # preparation to avoid duplicate computation.
90 + if self.mm_prefix_range is None:
91 + return None
92 +
93 + num_seqs = self.seq_lens.shape[0]
94 + device = self.seq_lens.device
95 +
96 + # Collect ranges, using [(0,0)] for empty sequences to ensure uniform dims
97 + range_lists = [
98 + self.mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
99 + ]
100 +
101 + # Return None if all ranges are trivial (only (0,0) placeholders)
102 + if all(r == [(0, 0)] for r in range_lists):
103 + return None
104 +
105 + # Create 2D tensors with shape (num_ranges, 2) for each sequence
106 + range_tensors = [
107 + torch.tensor(r, dtype=torch.int32, device=device).view(-1, 2)
108 + for r in range_lists
109 + ]
110 +
111 + return torch.nested.nested_tensor(range_tensors).to_padded_tensor(0)
112
113
114 class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
@@ -268,6 +301,10 @@ class TritonAttentionBackend(AttentionBackend):
301 def supports_head_size(cls, head_size: int) -> bool:
302 return head_size >= 32
303
304 + @classmethod
305 + def supports_mm_prefix(cls) -> bool:
306 + return True
307 +
308 @classmethod
309 def supports_sink(cls) -> bool:
310 return True
@@ -427,6 +464,7 @@ class TritonAttentionImpl(AttentionImpl):
464 softmax_segm_expsum = attn_metadata.softmax_segm_expsum
465
466 descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
467 + mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor
468
469 unified_attention(
470 q=query[:num_actual_tokens],
@@ -453,6 +491,7 @@ class TritonAttentionImpl(AttentionImpl):
491 softmax_segm_expsum=softmax_segm_expsum,
492 sinks=self.sinks,
493 output_scale=output_scale,
494 + mm_prefix_range=mm_prefix_range_tensor,
495 )
496
497 return output