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

Add Triton Backend for Unlimited-OCR R-SWA (#47102)

Browse Source 
Signed-off-by: Dakai An <dakaian108@gmail.com>
This commit is contained in:
Dakai An
2026-07-03 00:11:50 -07:00
committed by GitHub
parent b790c84cde
commit 1f486d96a1
5 changed files with 68 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/models/config.py
+14 -6
View File
@@ -71,8 +71,11 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'``
→ FlexAttention R-SWA via Triton block mask.
3. ``--attention-config '{"backend": "auto"}'`` (or omitted)
→ Auto-detect: FA4 if available (H20/H100 SM90), else FlexAttention.
3. ``--attention-config '{"backend": "TRITON_ATTN"}'``
→ Triton unified attention with an R-SWA decode mask.
4. ``--attention-config '{"backend": "auto"}'`` (or omitted)
→ Auto-detect: FA4 if available (H20/H100 SM90), else TritonAttention.
Regardless of backend, prefix caching is disabled for this model: R-SWA
decode-phase KV is not a pure causal function of the prefix (so decode
@@ -96,7 +99,7 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
attn_config.backend = (
AttentionBackendEnum.FLASH_ATTN
if fa4_available
else AttentionBackendEnum.FLEX_ATTENTION
else AttentionBackendEnum.TRITON_ATTN
)
logger.info(
"Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).",
@@ -110,8 +113,8 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
raise RuntimeError(
"Unlimited-OCR: --attention-config backend=FLASH_ATTN "
"requires FA4 (rswa_mask_mod), but FA4 is not available on "
"this device/installation. Use backend=FLEX_ATTENTION or "
"upgrade vllm-flash-attn."
"this device/installation. Use backend=TRITON_ATTN or "
"FLEX_ATTENTION, or upgrade vllm-flash-attn."
)
# On SM90 (H20), the default FA version is FA3 regardless of FA4
# availability (FA4 is only auto-upgraded when head_size > 256).
@@ -130,6 +133,11 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
attn_config.flash_attn_version,
)
elif attn_config.backend == AttentionBackendEnum.TRITON_ATTN:
logger.info(
"Unlimited-OCR: TritonAttention — R-SWA via unified attention mask."
)
elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION:
logger.info(
"Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.",
@@ -144,7 +152,7 @@ class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
raise ValueError(
f"Unlimited-OCR: unsupported attention backend "
f"{attn_config.backend!r} for R-SWA. "
"Use FLASH_ATTN (FA4) or FLEX_ATTENTION."
"Use FLASH_ATTN (FA4), TRITON_ATTN or FLEX_ATTENTION."
)
# R-SWA windows the *generated* tokens, so a decode-token's KV is not a
vllm/model_executor/models/unlimited_ocr.py
+5 -9
View File
@@ -21,15 +21,11 @@ backbone architecture, FlexAttention for R-SWA, vision encoder backend, and
Attention backend: the reference applies Reference Sliding Window Attention
(R-SWA) -- the prompt/image tokens form a globally-visible prefix while the
*generated* tokens additionally attend only a fixed sliding window (128) of
recent tokens. We reproduce this (Level 1: full KV cache + custom mask) by
forcing the language model onto the FlexAttention backend and installing an
R-SWA ``mask_mod``. FlexAttention is the only backend able to express the
"global prefix + sliding window" mask; FlashAttention-3 / Triton only support a
uniform window (and additionally crash or compute incorrectly on this decoder's
10-head,
head_dim-128 shape), and FlashInfer's paged decode exposes no custom mask. The
window size is published via ``model_config.rswa_window``, which the model
runner reads to plumb per-request prefix lengths into the FlexAttention mask.
recent tokens. We reproduce this with backend-specific custom masks: FA4 uses a
``mask_mod``, FlexAttention uses a Triton block mask, and TritonAttention uses a
unified-attention decode mask. FlashInfer's paged decode exposes no custom mask.
The window size is published via ``model_config.rswa_window``, which the model
runner reads to plumb per-request prefix lengths into the attention metadata.
The *vision encoder* (DeepEncoder's CLIP stage, head_dim 64) is unaffected and
does not use R-SWA: it runs a single full-attention prefill pass. FlashAttention,
vllm/v1/attention/backends/triton_attn.py
+23
View File
@@ -94,6 +94,8 @@ class TritonAttentionMetadata:
prefix_scheduler_metadata: torch.Tensor | None = None
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
mm_prefix_range_tensor: torch.Tensor | None = None
rswa_prefix_lens: torch.Tensor | None = None
rswa_window: int | None = None
class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
@@ -171,6 +173,14 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
dtype=torch.float32,
device=device,
)
self.rswa_window = model_config.rswa_window
self.persistent_rswa_prefix_lens: torch.Tensor | None = None
if self.rswa_window is not None:
self.persistent_rswa_prefix_lens = torch.empty(
vllm_config.scheduler_config.max_num_seqs,
dtype=torch.int32,
device=device,
)
def build_for_cudagraph_capture(
self, common_attn_metadata: CommonAttentionMetadata
@@ -244,6 +254,17 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
mm_ranges, num_reqs, seq_lens.device
)
rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens
if self.rswa_window is not None and rswa_prefix_lens is not None:
assert self.persistent_rswa_prefix_lens is not None
rswa_prefix_lens = rswa_prefix_lens.to(
device=self.device, dtype=torch.int32, non_blocking=True
)
persistent_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]
persistent_prefix_lens.copy_(rswa_prefix_lens[:num_reqs])
attn_metadata.rswa_prefix_lens = persistent_prefix_lens
attn_metadata.rswa_window = self.rswa_window
return attn_metadata
@@ -671,6 +692,8 @@ class TritonAttentionImpl(AttentionImpl):
sinks=self.sinks,
output_scale=output_scale,
mm_prefix_range=mm_prefix_range_tensor,
rswa_prefix_lens=attn_metadata.rswa_prefix_lens,
rswa_window=attn_metadata.rswa_window,
kv_quant_mode=self._kv_quant_mode,
k_scale_cache=k_scale_cache,
v_scale_cache=v_scale_cache,
vllm/v1/attention/ops/triton_attention_helpers.py
+10 -1
View File
@@ -281,6 +281,9 @@ def compute_kv_seq_mask(
USE_CAUSAL: tl.constexpr = True,
USE_PER_SEQ_CAUSAL: tl.constexpr = False,
per_seq_causal_ptr=None,
rswa_prefix_lens_ptr=None,
R_SWA_WINDOW: tl.constexpr = 0,
USE_R_SWA: tl.constexpr = False,
CHUNK_LOOKBACK: tl.constexpr = -1,
CHUNK_SIZE: tl.constexpr = -1,
MM_PREFIX_CLAMP_SW: tl.constexpr = False,
@@ -322,7 +325,7 @@ def compute_kv_seq_mask(
(query_abs_pos // CHUNK_SIZE - seq_offset[None, :] // CHUNK_SIZE)
<= CHUNK_LOOKBACK
)
elif SLIDING_WINDOW > 0:
elif SLIDING_WINDOW > 0 and not USE_R_SWA:
sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW
if USE_PER_SEQ_CAUSAL:
sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW
@@ -333,6 +336,12 @@ def compute_kv_seq_mask(
else:
seq_mask = seq_mask & sw_left
if USE_R_SWA:
prefix_len = tl.load(rswa_prefix_lens_ptr + seq_idx)
in_prefix = seq_offset[None, :] < prefix_len
in_window = (query_abs_pos - seq_offset) < R_SWA_WINDOW
seq_mask = seq_mask & (in_prefix | in_window)
# PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
# Default (MM_PREFIX_CLAMP_SW=False): applied AFTER sliding window so
# mm_prefix ranges override the SW restriction -> (causal AND SW) OR mm.
vllm/v1/attention/ops/triton_unified_attention.py
+16 -1
View File
@@ -221,6 +221,9 @@ def kernel_unified_attention(
USE_MM_PREFIX: tl.constexpr, # bool
MAX_MM_RANGES: tl.constexpr, # int
mm_prefix_range_ptr,
rswa_prefix_lens_ptr,
R_SWA_WINDOW: tl.constexpr, # int
USE_R_SWA: tl.constexpr, # bool
stride_k_cache_0: tl.int64, # int
stride_k_cache_1: tl.int64, # int
stride_k_cache_2: tl.int64, # int
@@ -395,7 +398,7 @@ def kernel_unified_attention(
BLOCK_Q,
num_queries_per_kv,
SLIDING_WINDOW,
USE_MM_PREFIX,
USE_MM_PREFIX or USE_R_SWA,
IS_3D,
USE_CAUSAL,
USE_PER_SEQ_CAUSAL,
@@ -511,6 +514,9 @@ def kernel_unified_attention(
USE_CAUSAL,
USE_PER_SEQ_CAUSAL,
per_seq_causal_ptr,
rswa_prefix_lens_ptr,
R_SWA_WINDOW,
USE_R_SWA,
CHUNK_LOOKBACK,
CHUNK_SIZE,
MM_PREFIX_CLAMP_SW,
@@ -812,6 +818,10 @@ def unified_attention(
sinks=None,
# Optional tensor for prefix lengths (PrefixLM support)
mm_prefix_range=None,
# R-SWA support: prefix tokens stay globally visible, generated tokens use
# a fixed sliding window.
rswa_prefix_lens=None,
rswa_window: int | None = None,
use_alibi_sqrt=False,
# KV cache quantization mode and per-token-head scale caches.
kv_quant_mode: KVQuantMode = KVQuantMode.NONE,
@@ -897,6 +907,8 @@ def unified_attention(
f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
)
use_rswa = rswa_window is not None and rswa_prefix_lens is not None
use_alibi_slopes = alibi_slopes is not None
use_qq_bias = qq_bias is not None
@@ -1111,6 +1123,9 @@ def unified_attention(
USE_MM_PREFIX=use_mm_prefix,
MAX_MM_RANGES=max_mm_ranges,
mm_prefix_range_ptr=mm_prefix_range,
rswa_prefix_lens_ptr=rswa_prefix_lens if use_rswa else seqused_k,
R_SWA_WINDOW=rswa_window or 0,
USE_R_SWA=use_rswa,
stride_k_cache_0=k.stride(0),
stride_k_cache_1=k.stride(1),
stride_k_cache_2=k.stride(2),