diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index a9e6815a49..cb01a12efd 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -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 diff --git a/vllm/model_executor/models/unlimited_ocr.py b/vllm/model_executor/models/unlimited_ocr.py index 06dc02512f..0f29337314 100644 --- a/vllm/model_executor/models/unlimited_ocr.py +++ b/vllm/model_executor/models/unlimited_ocr.py @@ -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, diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index b0030d27d9..6ba0c94537 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -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, diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index e753d53cdf..b2667935c6 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -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. diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index a1c9abde88..a13ae96a7a 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -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),