Add Triton Backend for Unlimited-OCR R-SWA (#47102)
Signed-off-by: Dakai An <dakaian108@gmail.com>
Dakai An committed
Jul 3, 2026 at 15:11 UTC
1f486d96a17303ce8db8e02be39545b2be338446
5 files changed
+68
-17
vllm/model_executor/models/config.py
+14
-6
@@ -71,8 +71,11 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
71
2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'``
72
→ FlexAttention R-SWA via Triton block mask.
73
74
- 3. ``--attention-config '{"backend": "auto"}'`` (or omitted)
75
- → Auto-detect: FA4 if available (H20/H100 SM90), else FlexAttention.
74
+ 3. ``--attention-config '{"backend": "TRITON_ATTN"}'``
75
+ → Triton unified attention with an R-SWA decode mask.
76
+
77
+ 4. ``--attention-config '{"backend": "auto"}'`` (or omitted)
78
+ → Auto-detect: FA4 if available (H20/H100 SM90), else TritonAttention.
79
80
Regardless of backend, prefix caching is disabled for this model: R-SWA
81
decode-phase KV is not a pure causal function of the prefix (so decode
@@ -96,7 +99,7 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
99
attn_config.backend = (
100
AttentionBackendEnum.FLASH_ATTN
101
if fa4_available
99
- else AttentionBackendEnum.FLEX_ATTENTION
102
+ else AttentionBackendEnum.TRITON_ATTN
103
)
104
logger.info(
105
"Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).",
@@ -110,8 +113,8 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
113
raise RuntimeError(
114
"Unlimited-OCR: --attention-config backend=FLASH_ATTN "
115
"requires FA4 (rswa_mask_mod), but FA4 is not available on "
113
- "this device/installation. Use backend=FLEX_ATTENTION or "
114
- "upgrade vllm-flash-attn."
116
+ "this device/installation. Use backend=TRITON_ATTN or "
117
+ "FLEX_ATTENTION, or upgrade vllm-flash-attn."
118
)
119
# On SM90 (H20), the default FA version is FA3 regardless of FA4
120
# availability (FA4 is only auto-upgraded when head_size > 256).
@@ -130,6 +133,11 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
133
attn_config.flash_attn_version,
134
)
135
136
+ elif attn_config.backend == AttentionBackendEnum.TRITON_ATTN:
137
+ logger.info(
138
+ "Unlimited-OCR: TritonAttention — R-SWA via unified attention mask."
139
+ )
140
+
141
elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION:
142
logger.info(
143
"Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.",
@@ -144,7 +152,7 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
152
raise ValueError(
153
f"Unlimited-OCR: unsupported attention backend "
154
f"{attn_config.backend!r} for R-SWA. "
147
- "Use FLASH_ATTN (FA4) or FLEX_ATTENTION."
155
+ "Use FLASH_ATTN (FA4), TRITON_ATTN or FLEX_ATTENTION."
156
)
157
158
# R-SWA windows the *generated* tokens, so a decode-token's KV is not a
vllm/model_executor/models/unlimited_ocr.py
+5
-9
@@ -21,15 +21,11 @@ backbone architecture, FlexAttention for R-SWA, vision encoder backend, and
21
Attention backend: the reference applies Reference Sliding Window Attention
22
(R-SWA) -- the prompt/image tokens form a globally-visible prefix while the
23
*generated* tokens additionally attend only a fixed sliding window (128) of
24
-recent tokens. We reproduce this (Level 1: full KV cache + custom mask) by
25
-forcing the language model onto the FlexAttention backend and installing an
26
-R-SWA ``mask_mod``. FlexAttention is the only backend able to express the
27
-"global prefix + sliding window" mask; FlashAttention-3 / Triton only support a
28
-uniform window (and additionally crash or compute incorrectly on this decoder's
29
-10-head,
30
-head_dim-128 shape), and FlashInfer's paged decode exposes no custom mask. The
31
-window size is published via ``model_config.rswa_window``, which the model
32
-runner reads to plumb per-request prefix lengths into the FlexAttention mask.
24
+recent tokens. We reproduce this with backend-specific custom masks: FA4 uses a
25
+``mask_mod``, FlexAttention uses a Triton block mask, and TritonAttention uses a
26
+unified-attention decode mask. FlashInfer's paged decode exposes no custom mask.
27
+The window size is published via ``model_config.rswa_window``, which the model
28
+runner reads to plumb per-request prefix lengths into the attention metadata.
29
30
The *vision encoder* (DeepEncoder's CLIP stage, head_dim 64) is unaffected and
31
does not use R-SWA: it runs a single full-attention prefill pass. FlashAttention,
vllm/v1/attention/backends/triton_attn.py
+23
@@ -94,6 +94,8 @@ class TritonAttentionMetadata:
94
prefix_scheduler_metadata: torch.Tensor | None = None
95
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
96
mm_prefix_range_tensor: torch.Tensor | None = None
97
+ rswa_prefix_lens: torch.Tensor | None = None
98
+ rswa_window: int | None = None
99
100
101
class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
@@ -171,6 +173,14 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
173
dtype=torch.float32,
174
device=device,
175
)
176
+ self.rswa_window = model_config.rswa_window
177
+ self.persistent_rswa_prefix_lens: torch.Tensor | None = None
178
+ if self.rswa_window is not None:
179
+ self.persistent_rswa_prefix_lens = torch.empty(
180
+ vllm_config.scheduler_config.max_num_seqs,
181
+ dtype=torch.int32,
182
+ device=device,
183
+ )
184
185
def build_for_cudagraph_capture(
186
self, common_attn_metadata: CommonAttentionMetadata
@@ -244,6 +254,17 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
254
mm_ranges, num_reqs, seq_lens.device
255
)
256
257
+ rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens
258
+ if self.rswa_window is not None and rswa_prefix_lens is not None:
259
+ assert self.persistent_rswa_prefix_lens is not None
260
+ rswa_prefix_lens = rswa_prefix_lens.to(
261
+ device=self.device, dtype=torch.int32, non_blocking=True
262
+ )
263
+ persistent_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]
264
+ persistent_prefix_lens.copy_(rswa_prefix_lens[:num_reqs])
265
+ attn_metadata.rswa_prefix_lens = persistent_prefix_lens
266
+ attn_metadata.rswa_window = self.rswa_window
267
+
268
return attn_metadata
269
270
@@ -671,6 +692,8 @@ class TritonAttentionImpl(AttentionImpl):
692
sinks=self.sinks,
693
output_scale=output_scale,
694
mm_prefix_range=mm_prefix_range_tensor,
695
+ rswa_prefix_lens=attn_metadata.rswa_prefix_lens,
696
+ rswa_window=attn_metadata.rswa_window,
697
kv_quant_mode=self._kv_quant_mode,
698
k_scale_cache=k_scale_cache,
699
v_scale_cache=v_scale_cache,
vllm/v1/attention/ops/triton_attention_helpers.py
+10
-1
@@ -281,6 +281,9 @@ def compute_kv_seq_mask(
281
USE_CAUSAL: tl.constexpr = True,
282
USE_PER_SEQ_CAUSAL: tl.constexpr = False,
283
per_seq_causal_ptr=None,
284
+ rswa_prefix_lens_ptr=None,
285
+ R_SWA_WINDOW: tl.constexpr = 0,
286
+ USE_R_SWA: tl.constexpr = False,
287
CHUNK_LOOKBACK: tl.constexpr = -1,
288
CHUNK_SIZE: tl.constexpr = -1,
289
MM_PREFIX_CLAMP_SW: tl.constexpr = False,
@@ -322,7 +325,7 @@ def compute_kv_seq_mask(
325
(query_abs_pos // CHUNK_SIZE - seq_offset[None, :] // CHUNK_SIZE)
326
<= CHUNK_LOOKBACK
327
)
325
- elif SLIDING_WINDOW > 0:
328
+ elif SLIDING_WINDOW > 0 and not USE_R_SWA:
329
sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW
330
if USE_PER_SEQ_CAUSAL:
331
sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW
@@ -333,6 +336,12 @@ def compute_kv_seq_mask(
336
else:
337
seq_mask = seq_mask & sw_left
338
339
+ if USE_R_SWA:
340
+ prefix_len = tl.load(rswa_prefix_lens_ptr + seq_idx)
341
+ in_prefix = seq_offset[None, :] < prefix_len
342
+ in_window = (query_abs_pos - seq_offset) < R_SWA_WINDOW
343
+ seq_mask = seq_mask & (in_prefix | in_window)
344
+
345
# PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
346
# Default (MM_PREFIX_CLAMP_SW=False): applied AFTER sliding window so
347
# mm_prefix ranges override the SW restriction -> (causal AND SW) OR mm.
vllm/v1/attention/ops/triton_unified_attention.py
+16
-1
@@ -221,6 +221,9 @@ def kernel_unified_attention(
221
USE_MM_PREFIX: tl.constexpr, # bool
222
MAX_MM_RANGES: tl.constexpr, # int
223
mm_prefix_range_ptr,
224
+ rswa_prefix_lens_ptr,
225
+ R_SWA_WINDOW: tl.constexpr, # int
226
+ USE_R_SWA: tl.constexpr, # bool
227
stride_k_cache_0: tl.int64, # int
228
stride_k_cache_1: tl.int64, # int
229
stride_k_cache_2: tl.int64, # int
@@ -395,7 +398,7 @@ def kernel_unified_attention(
398
BLOCK_Q,
399
num_queries_per_kv,
400
SLIDING_WINDOW,
398
- USE_MM_PREFIX,
401
+ USE_MM_PREFIX or USE_R_SWA,
402
IS_3D,
403
USE_CAUSAL,
404
USE_PER_SEQ_CAUSAL,
@@ -511,6 +514,9 @@ def kernel_unified_attention(
514
USE_CAUSAL,
515
USE_PER_SEQ_CAUSAL,
516
per_seq_causal_ptr,
517
+ rswa_prefix_lens_ptr,
518
+ R_SWA_WINDOW,
519
+ USE_R_SWA,
520
CHUNK_LOOKBACK,
521
CHUNK_SIZE,
522
MM_PREFIX_CLAMP_SW,
@@ -812,6 +818,10 @@ def unified_attention(
818
sinks=None,
819
# Optional tensor for prefix lengths (PrefixLM support)
820
mm_prefix_range=None,
821
+ # R-SWA support: prefix tokens stay globally visible, generated tokens use
822
+ # a fixed sliding window.
823
+ rswa_prefix_lens=None,
824
+ rswa_window: int | None = None,
825
use_alibi_sqrt=False,
826
# KV cache quantization mode and per-token-head scale caches.
827
kv_quant_mode: KVQuantMode = KVQuantMode.NONE,
@@ -897,6 +907,8 @@ def unified_attention(
907
f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
908
)
909
910
+ use_rswa = rswa_window is not None and rswa_prefix_lens is not None
911
+
912
use_alibi_slopes = alibi_slopes is not None
913
use_qq_bias = qq_bias is not None
914
@@ -1111,6 +1123,9 @@ def unified_attention(
1123
USE_MM_PREFIX=use_mm_prefix,
1124
MAX_MM_RANGES=max_mm_ranges,
1125
mm_prefix_range_ptr=mm_prefix_range,
1126
+ rswa_prefix_lens_ptr=rswa_prefix_lens if use_rswa else seqused_k,
1127
+ R_SWA_WINDOW=rswa_window or 0,
1128
+ USE_R_SWA=use_rswa,
1129
stride_k_cache_0=k.stride(0),
1130
stride_k_cache_1=k.stride(1),
1131
stride_k_cache_2=k.stride(2),