[Model] Support Unlimited OCR

docs/models/supported_models.md

@@ -629,6 +629,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
+| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I<sup>+</sup> | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ |
 
 Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it!
 

tests/models/registry.py

@@ -825,6 +825,9 @@ def check_available_online(
     "DeepseekOCR2ForCausalLM": _HfExamplesInfo(
         "deepseek-ai/DeepSeek-OCR-2",
     ),
+    "UnlimitedOCRForCausalLM": _HfExamplesInfo(
+        "baidu/Unlimited-OCR",
+    ),
     "DotsOCRForCausalLM": _HfExamplesInfo(
         "rednote-hilab/dots.ocr", trust_remote_code=True
     ),

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -14,9 +14,14 @@
 )
 from vllm.v1.core.single_type_kv_cache_manager import (
     ChunkedLocalAttentionManager,
+    RSWAManager,
     SlidingWindowManager,
 )
-from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec
+from vllm.v1.kv_cache_interface import (
+    ChunkedLocalAttentionSpec,
+    RSWASpec,
+    SlidingWindowSpec,
+)
 
 pytestmark = pytest.mark.cpu_test
 
@@ -327,6 +332,51 @@ def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
     assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
 
 
+def test_rswa_remove_skipped_blocks_gap_range():
+    block_size = 4
+    rswa_spec = RSWASpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        rswa_window=8,
+    )
+    block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=4)
+    manager = RSWAManager(
+        rswa_spec,
+        block_pool=block_pool,
+        enable_caching=True,
+        kv_cache_group_id=0,
+        scheduler_block_size=block_size,
+    )
+
+    null_block_id = block_pool.null_block.block_id
+    original_block_ids = list(range(1000, 1010))
+    block_table = [
+        KVCacheBlock(id_) if id_ != null_block_id else block_pool.null_block
+        for id_ in original_block_ids
+    ]
+    manager.req_to_blocks["test"] = block_table
+
+    prefix_len = 16
+
+    # Without num_prompt_tokens, R-SWA does not evict gap blocks.
+    manager.remove_skipped_blocks("test", 28)
+    assert [b.block_id for b in block_table] == original_block_ids
+
+    # Gap = block 4 only (tokens [16, 20) fall in the gap).
+    manager.remove_skipped_blocks("test", 28, num_prompt_tokens=prefix_len)
+    expected = original_block_ids.copy()
+    expected[4] = null_block_id
+    assert [b.block_id for b in block_table] == expected
+
+    # Window moves: blocks 5 and 6 also enter the gap; block 4 is already null.
+    manager.remove_skipped_blocks("test", 36, num_prompt_tokens=prefix_len)
+    expected[5] = null_block_id
+    expected[6] = null_block_id
+    assert [b.block_id for b in block_table] == expected
+
+
 def test_get_num_blocks_to_allocate():
     block_size = 2
     sliding_window_spec = SlidingWindowSpec(

vllm/config/model.py

@@ -1251,6 +1251,10 @@ def is_deepseek_mla(self) -> bool:
     def is_mm_prefix_lm(self) -> bool:
         return self.model_arch_config.is_mm_prefix_lm
 
+    @property
+    def rswa_window(self) -> int | None:
+        return self.model_arch_config.rswa_window
+
     def get_head_size(self) -> int:
         return self.model_arch_config.head_size
 

vllm/config/model_arch.py

@@ -56,5 +56,8 @@ class ModelArchitectureConfig:
     is_mm_prefix_lm: bool
     """Whether the model uses image bidirectional attention."""
 
+    rswa_window: int | None
+    """Reference Sliding Window Attention window size (None disables R-SWA)."""
+
     derived_max_model_len_and_key: tuple[float, str | None]
     """Derived maximum model length and key from the hf config."""

vllm/model_executor/layers/attention/__init__.py

@@ -14,6 +14,7 @@
 from vllm.model_executor.layers.attention.prefill_prefix_lm_attention import (
     PrefillPrefixLMAttention,
 )
+from vllm.model_executor.layers.attention.rswa_attention import RSWAAttention
 from vllm.model_executor.layers.attention.static_sink_attention import (
     StaticSinkAttention,
 )
@@ -26,5 +27,6 @@
     "MLAAttention",
     "MMEncoderAttention",
     "PrefillPrefixLMAttention",
+    "RSWAAttention",
     "StaticSinkAttention",
 ]

vllm/model_executor/layers/attention/rswa_attention.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.config.vllm import VllmConfig
+from vllm.model_executor.layers.attention import Attention
+from vllm.v1.kv_cache_interface import KVCacheSpec, RSWASpec, get_kv_quant_mode
+
+
+class RSWAAttention(Attention):
+    """Attention layer that reports ``RSWASpec`` as its KV cache spec.
+
+    Drop-in replacement for the standard ``Attention`` layer when the model is
+    configured with Reference Sliding Window Attention (R-SWA,
+    ``rswa_window > 0``). The actual masking logic lives in the attention
+    backend (FlexAttention or FA4 mask_mod); this layer only overrides
+    ``get_kv_cache_spec`` so the KV cache manager instantiates ``RSWAManager``
+    (instead of ``FullAttentionManager``) and can therefore evict "gap" blocks
+    to keep per-request KV memory bounded at O(prefix + window).
+    """
+
+    def __init__(self, *args, rswa_window: int, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._rswa_window = rswa_window
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+        spec = super().get_kv_cache_spec(vllm_config)
+        if spec is None:
+            return None
+        return RSWASpec(
+            block_size=vllm_config.cache_config.block_size,
+            num_kv_heads=self.num_kv_heads,
+            head_size=self.head_size,
+            head_size_v=self.head_size_v,
+            dtype=self.kv_cache_torch_dtype,
+            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
+            rswa_window=self._rswa_window,
+        )

vllm/model_executor/models/config.py

@@ -54,6 +54,139 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         hf_config.is_causal = not hf_config.use_bidirectional_attention
 
 
+class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        """Configure Unlimited-OCR attention backends for R-SWA and vision.
+
+        Backend selection — controlled by the standard ``--attention-config``
+        CLI argument (priority order):
+
+          1. ``--attention-config '{"backend": "FLASH_ATTN"}'``
+             → FA4 + rswa_mask_mod.  Exact token-level R-SWA.
+               ``flash_attn_version`` is forced to 4 if not already set (R-SWA
+               mask_mod requires FA4; FA3 cannot express it).  Raises if FA4 is
+               not available on this device.
+
+          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.
+
+        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
+        blocks are not reusable), and single-turn image-led OCR prompts rarely
+        hit the prefix cache.
+
+        Example — force FlexAttention even on a machine with FA4::
+
+            vllm serve baidu/Unlimited-OCR \\
+                --attention-config '{"backend": "FLEX_ATTENTION"}'
+        """
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+        from vllm.vllm_flash_attn import is_fa_version_supported
+
+        attn_config = vllm_config.attention_config
+        fa4_available = is_fa_version_supported(4)
+
+        # ── step 1: resolve backend ─────────────────────────────────────────
+        # None means the user did not explicitly specify a backend; auto-select.
+        if attn_config.backend is None:
+            attn_config.backend = (
+                AttentionBackendEnum.FLASH_ATTN
+                if fa4_available
+                else AttentionBackendEnum.FLEX_ATTENTION
+            )
+            logger.info(
+                "Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).",
+                attn_config.backend.value,
+                fa4_available,
+            )
+
+        # ── step 2: configure the chosen backend ────────────────────────────
+        if attn_config.backend == AttentionBackendEnum.FLASH_ATTN:
+            if not fa4_available:
+                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."
+                )
+            # On SM90 (H20), the default FA version is FA3 regardless of FA4
+            # availability (FA4 is only auto-upgraded when head_size > 256).
+            # The R-SWA mask_mod requires FA4, so force the version globally.
+            if attn_config.flash_attn_version is None:
+                attn_config.flash_attn_version = 4
+            elif attn_config.flash_attn_version < 4:
+                logger.warning(
+                    "Unlimited-OCR: flash_attn_version=%d cannot express the "
+                    "R-SWA mask_mod; upgrading to 4.",
+                    attn_config.flash_attn_version,
+                )
+                attn_config.flash_attn_version = 4
+            logger.info(
+                "Unlimited-OCR: FlashAttention FA%d + rswa_mask_mod — exact R-SWA.",
+                attn_config.flash_attn_version,
+            )
+
+        elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION:
+            logger.info(
+                "Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.",
+                ""
+                if not fa4_available
+                else (
+                    " (FA4 available but not used; pass backend=FLASH_ATTN to upgrade)"
+                ),
+            )
+
+        else:
+            raise ValueError(
+                f"Unlimited-OCR: unsupported attention backend "
+                f"{attn_config.backend!r} for R-SWA. "
+                "Use FLASH_ATTN (FA4) or FLEX_ATTENTION."
+            )
+
+        # R-SWA windows the *generated* tokens, so a decode-token's KV is not a
+        # pure causal function of the prefix and cannot be safely reused across
+        # requests via prefix caching. Only the prompt/image prefix is cacheable,
+        # but OCR is single-turn with image-led prompts that rarely share a
+        # prefix, so prefix caching brings little benefit while complicating the
+        # KV cache manager. Disable it for this model.
+        cache_config = vllm_config.cache_config
+        if cache_config.enable_prefix_caching:
+            cache_config.enable_prefix_caching = False
+            logger.info(
+                "Unlimited-OCR: disabling prefix caching (R-SWA decode KV is not "
+                "cacheable, and single-turn image-led prompts rarely hit the "
+                "prefix cache)."
+            )
+
+        mm_config = getattr(vllm_config.model_config, "multimodal_config", None)
+        if mm_config is not None:
+            if mm_config.mm_encoder_attn_backend is None:
+                mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN
+            elif mm_config.mm_encoder_attn_backend == AttentionBackendEnum.FLASHINFER:
+                logger.warning(
+                    "Unlimited-OCR: FlashInfer is not supported for the vision "
+                    "encoder (the CLIP stage runs full attention without "
+                    "cu_seqlens); falling back to FlashAttention."
+                )
+                mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN
+
+    @staticmethod
+    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
+        text_config = model_config.hf_config.text_config
+        text_config.architectures = ["DeepseekV2ForCausalLM"]
+        if getattr(model_config.hf_config, "rswa_window", None) is None:
+            model_config.hf_config.rswa_window = 128
+        # Propagate rswa_window to text_config so that DeepseekAttention (which
+        # receives text_config as its vllm_config.model_config.hf_config via
+        # init_vllm_registered_model) can read it and create RSWAAttention.
+        rswa_window = model_config.hf_config.rswa_window
+        text_config.rswa_window = rswa_window
+
+
 class Gemma4Config(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -703,6 +836,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
     "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
     "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
     "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
+    "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig,
     "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
     "XLMRobertaModel": JinaRobertaModelConfig,
 }

vllm/model_executor/models/deepseek_ocr.py

@@ -220,8 +220,10 @@ def get_num_image_tokens(
         patch_size = 16
         downsample_ratio = 4
 
-        if CROP_MODE:
-            if image_width <= 640 and image_height <= 640:
+        # Use the caller-supplied `cropping` flag so that callers that disable
+        # crop mode for multi-image requests get a consistent token count.
+        if cropping:
+            if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE:
                 crop_ratio = [1, 1]
             else:
                 # find the closest aspect ratio to the target

vllm/model_executor/models/deepseek_v2.py

@@ -45,7 +45,7 @@
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention import Attention, RSWAAttention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
@@ -174,15 +174,28 @@ def __init__(
             max_position=max_position_embeddings,
             rope_parameters=config.rope_parameters,
         )
-        self.attn = Attention(
-            self.num_heads,
-            self.head_dim,
-            self.scaling,
-            num_kv_heads=self.num_kv_heads,
-            cache_config=cache_config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.attn",
-        )
+        rswa_window = getattr(vllm_config.model_config.hf_config, "rswa_window", None)
+        if rswa_window is not None:
+            self.attn = RSWAAttention(
+                self.num_heads,
+                self.head_dim,
+                self.scaling,
+                num_kv_heads=self.num_kv_heads,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.attn",
+                rswa_window=rswa_window,
+            )
+        else:
+            self.attn = Attention(
+                self.num_heads,
+                self.head_dim,
+                self.scaling,
+                num_kv_heads=self.num_kv_heads,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.attn",
+            )
 
     def forward(
         self,
@@ -588,12 +601,12 @@ def __init__(
         compilation_config.static_forward_context[prefix] = self
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
-        return MLAAttentionSpec(  # Only has one vector instead of K + V
+        return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
-        )
+        )  # Only has one vector instead of K + V
 
     def forward(self): ...
 

vllm/model_executor/models/registry.py

@@ -358,6 +358,7 @@
     "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
     "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
     "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
+    "UnlimitedOCRForCausalLM": ("unlimited_ocr", "UnlimitedOCRForCausalLM"),
     "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
     "Eagle2_5_VLForConditionalGeneration": (
         "eagle2_5_vl",