[Model] Support Unlimited OCR (#46564)

Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Roger Wang <hey@rogerw.io>

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 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "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.kv_cache_utils import (
 )
 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 test_sliding_window_remove_skipped_blocks():
     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 @@ class ModelConfig:
     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.mm_encoder_attention import MMEncoderA
 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 @@ __all__ = [
     "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 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
         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 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
     "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 @@ class DeepseekOCRProcessingInfo(BaseProcessingInfo):
         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.distributed import (
 )
 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 @@ class DeepseekAttention(nn.Module):
             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 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
         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 @@ _MULTIMODAL_MODELS = {
     "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",

vllm/model_executor/models/unlimited_ocr.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Unlimited-OCR model compatible with HuggingFace weights.
+
+Unlimited-OCR (``baidu/Unlimited-OCR``) shares
+the exact DeepSeek-OCR (gundam, ``base_size=1024`` / ``image_size=640`` / crop)
+vision stack: a DeepEncoder (SAM-ViT-B + CLIP-L) followed by a linear MLP
+projector, with the same image-token tiling layout. The only difference is the
+language backbone, which is a DeepSeek-V2 *MoE* (64 routed + 2 shared experts,
+``first_k_dense_replace=1``) that uses plain multi-head attention
+(``use_mla=False``, ``qk_nope_head_dim == qk_rope_head_dim == 0``) instead of
+the dense MLA decoder used by DeepSeek-OCR.
+
+vLLM's ``DeepseekV2DecoderLayer`` already dispatches to the plain-MHA
+``DeepseekAttention`` whenever ``qk_nope_head_dim == qk_rope_head_dim == 0`` and
+builds the MoE blocks straight from the config, so the whole DeepSeek-OCR
+multimodal wrapper can be reused verbatim. Model-specific config (language
+backbone architecture, FlexAttention for R-SWA, vision encoder backend, and
+``rswa_window``) is applied in ``UnlimitedOCRForCausalLMConfig``.
+
+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.
+
+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,
+Triton and torch SDPA all produce correct, equally fast results; only FlashInfer
+is incompatible (its ViT path asserts on the varlen cu_seqlens metadata that this
+CLIP encoder never builds). We default the encoder to FlashAttention and
+transparently fall back to it if FlashInfer is requested.
+
+To suppress repetition on long documents, use ``NGramPerReqLogitsProcessor`` from
+this module (same request-level processor as DeepSeek-OCR) with::
+
+    SamplingParams(
+        temperature=0.0,
+        max_tokens=8192,
+        extra_args={"ngram_size": 35, "window_size": 128},
+    )
+
+Image processing
+----------------
+Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR), i.e.
+``dynamic_preprocess`` runs with ``max_num=32``.
+
+Multi-image requests fall back to non-crop mode: crop ("gundam") mode is only
+used for single-image input. DeepSeek-OCR does *not* have this restriction.
+
+Because that fallback makes the per-image processor output depend on *how many*
+images are in the request, it breaks the assumption behind vLLM's per-item
+multimodal processing cache (``MultiModalProcessorOnlyCache``). We handle this
+the same way ``DeepseekVL2MultiModalProcessor`` does: only the single-image case
+(which always crops) is cached, while multi-image requests bypass the cache and
+are recomputed fresh -- see ``_cached_apply_hf_processor`` below. This keeps the
+processing cache consistent (verified by ``test_processing_correctness``).
+"""
+
+import math
+from collections.abc import Mapping, Sequence
+
+from vllm.config import VllmConfig
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalKwargsItems
+from vllm.multimodal.parse import (
+    ImageEmbeddingItems,
+    ImageProcessorItems,
+    ImageSize,
+    MultiModalDataItems,
+)
+from vllm.multimodal.processing import PromptReplacement, PromptUpdate
+from vllm.multimodal.processing.context import TimingContext
+from vllm.multimodal.processing.inputs import ProcessorInputs
+from vllm.multimodal.processing.processor import MultiModalProcessingInfo
+from vllm.transformers_utils.processors.deepseek_ocr import (
+    BASE_SIZE,
+    CROP_MODE,
+    IMAGE_SIZE,
+    count_tiles,
+)
+
+from .deepseek_ocr import (
+    DeepseekOCRDummyInputsBuilder,
+    DeepseekOCRForCausalLM,
+    DeepseekOCRMultiModalProcessor,
+    DeepseekOCRProcessingInfo,
+    NGramPerReqLogitsProcessor,
+)
+
+__all__ = [
+    "NGramPerReqLogitsProcessor",
+    "UnlimitedOCRForCausalLM",
+]
+
+# Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR).
+_UNLIMITED_OCR_MAX_CROPS = 32
+
+
+class UnlimitedOCRProcessingInfo(DeepseekOCRProcessingInfo):
+    """ProcessingInfo for Unlimited-OCR: same as DeepSeek-OCR but with
+    max_crops=32 instead of 6.  The higher crop count allows tiling very large
+    document pages into up to 32 640×640 patches (dynamic_preprocess max_num=32).
+    """
+
+    def get_hf_config(self):
+        from vllm.transformers_utils.configs.unlimited_ocr import UnlimitedOCRConfig
+
+        return self.ctx.get_hf_config(UnlimitedOCRConfig)
+
+    def get_hf_processor(self, **kwargs: object):
+        from vllm.transformers_utils.processors.unlimited_ocr import (
+            UnlimitedOCRProcessor,
+        )
+
+        v1_processor_config = dict(
+            image_size=IMAGE_SIZE,
+            base_size=BASE_SIZE,
+            crop_mode=CROP_MODE,
+            strategy="v1",
+            max_crops=_UNLIMITED_OCR_MAX_CROPS,
+        )
+        return self.ctx.get_hf_processor(
+            UnlimitedOCRProcessor,
+            **{**v1_processor_config, **kwargs},
+        )
+
+    def get_num_image_tokens(
+        self, *, image_width: int, image_height: int, cropping: bool = True
+    ) -> int:
+        patch_size = 16
+        downsample_ratio = 4
+
+        # Honour the caller-supplied `cropping` flag: multi-image callers pass
+        # cropping=False to match UnlimitedOCRProcessor.tokenize_with_images.
+        if cropping:
+            if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE:
+                crop_ratio = [1, 1]
+            else:
+                crop_ratio = count_tiles(
+                    image_width,
+                    image_height,
+                    max_num=_UNLIMITED_OCR_MAX_CROPS,
+                    image_size=IMAGE_SIZE,
+                )
+            num_width_tiles, num_height_tiles = crop_ratio
+        else:
+            num_width_tiles = num_height_tiles = 1
+
+        h = w = math.ceil((BASE_SIZE // patch_size) / downsample_ratio)
+        h2 = w2 = math.ceil((IMAGE_SIZE // patch_size) / downsample_ratio)
+
+        global_views_tokens = h * (w + 1)
+        if num_width_tiles > 1 or num_height_tiles > 1:
+            local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1)
+        else:
+            local_views_tokens = 0
+
+        return global_views_tokens + local_views_tokens + 1
+
+    def get_image_size_with_most_features(self) -> ImageSize:
+        # With max_crops=32, the widest possible grid is 4×8 (aspect ratio 1:2).
+        # A 2560×5120 image (4×640 × 8×640) selects exactly 4×8=32 tiles and
+        # produces the maximum token count.
+        return ImageSize(width=640 * 4, height=640 * 8)
+
+
+class UnlimitedOCRMultiModalProcessor(DeepseekOCRMultiModalProcessor):
+    """Multimodal processor for Unlimited-OCR.
+
+    Disables crop mode for multi-image requests (to stay consistent with
+    ``UnlimitedOCRProcessor.tokenize_with_images``), and -- since that makes the
+    per-image output depend on the request's image count -- bypasses the
+    per-item processing cache for multi-image requests, exactly like
+    ``DeepseekVL2MultiModalProcessor``.
+
+    DeepSeek-OCR does *not* apply either of these.
+    """
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        image_token_id = hf_processor.image_token_id
+        assert isinstance(image_token_id, int)
+
+        def get_replacement_unlimited_ocr(item_idx: int):
+            images = mm_items.get_items(
+                "image", (ImageEmbeddingItems, ImageProcessorItems)
+            )
+
+            if isinstance(images, ImageEmbeddingItems):
+                num_image_tokens = images.get_feature_size(item_idx)
+            else:
+                size = images.get_image_size(item_idx)
+
+                # Disable crop mode for multi-image input.
+                # UnlimitedOCRProcessor.tokenize_with_images applies the same
+                # fallback, so both paths must agree on the effective crop flag.
+                effective_cropping = CROP_MODE and len(images) == 1
+
+                num_image_tokens = self.info.get_num_image_tokens(
+                    image_width=size.width,
+                    image_height=size.height,
+                    cropping=effective_cropping,
+                )
+            return [image_token_id] * num_image_tokens
+
+        return [
+            PromptReplacement(
+                modality="image",
+                target=[image_token_id],
+                replacement=get_replacement_unlimited_ocr,
+            )
+        ]
+
+    def _cached_apply_hf_processor(
+        self,
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
+        # The processor logic differs for single-image (crop) vs multi-image
+        # (no crop) requests. The processing cache assumes per-item output is
+        # invariant of how many images are passed per prompt, so we only cache
+        # the single-image case and recompute multi-image requests fresh.
+        if inputs.mm_data_items.get_count("image", strict=False) > 1:
+            return self._apply_hf_processor(inputs, timing_ctx)
+
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    UnlimitedOCRMultiModalProcessor,
+    info=UnlimitedOCRProcessingInfo,
+    dummy_inputs=DeepseekOCRDummyInputsBuilder,
+)
+class UnlimitedOCRForCausalLM(DeepseekOCRForCausalLM):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)

vllm/tokenizers/registry.py

@@ -31,7 +31,11 @@ logger = init_logger(__name__)
 # temporary workaround and better long term solutions are:
 # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better)
 # - Fix tokenizer_class on the hub for the affected models (best)
-_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"}
+_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {
+    "step3_vl",
+    "step3p7",
+    "unlimited-ocr",
+}
 
 _VLLM_TOKENIZERS = {
     "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),

vllm/transformers_utils/chat_templates/registry.py

@@ -29,6 +29,7 @@ _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: dict[str, ChatTemplatePath] = {
     "colpali": CHAT_TEMPLATES_DIR / "template_basic.jinja",
     "deepseek_ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",
     "deepseek_ocr2": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",
+    "unlimited-ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",
     "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja",
     "fuyu": CHAT_TEMPLATES_DIR / "template_fuyu.jinja",
     "minicpmv": _get_minicpmv_chat_template_fallback,

vllm/transformers_utils/config.py

@@ -124,6 +124,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     laguna="LagunaConfig",
     lfm2_moe="Lfm2MoeConfig",
     tarsier2="Tarsier2Config",
+    **{"unlimited-ocr": "UnlimitedOCRConfig"},
 )
 
 _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"}

vllm/transformers_utils/configs/__init__.py

@@ -73,6 +73,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "RadioConfig": "vllm.transformers_utils.configs.radio",
     "SpeculatorsConfig": "vllm.transformers_utils.configs.speculators",
     "UltravoxConfig": "vllm.transformers_utils.configs.ultravox",
+    "UnlimitedOCRConfig": "vllm.transformers_utils.configs.unlimited_ocr",
     "Step3VLConfig": "vllm.transformers_utils.configs.step3_vl",
     "Step3VisionEncoderConfig": "vllm.transformers_utils.configs.step3_vl",
     "Step3TextConfig": "vllm.transformers_utils.configs.step3_vl",
@@ -147,6 +148,7 @@ __all__ = [
     "RadioConfig",
     "SpeculatorsConfig",
     "UltravoxConfig",
+    "UnlimitedOCRConfig",
     "Step3VLConfig",
     "Step3VisionEncoderConfig",
     "Step3TextConfig",

vllm/transformers_utils/configs/unlimited_ocr.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Unlimited-OCR (baidu/Unlimited-OCR) reuses
+# the DeepSeek-OCR multimodal layout (DeepEncoder = SAM-ViT-B + CLIP-L, a linear
+# MLP projector and a DeepSeek-V2 text backbone). The only architectural
+# difference is the language model, which is a DeepSeek-V2 *MoE* with plain
+# multi-head attention (``use_mla=False``) instead of the dense MLA backbone.
+# We therefore reuse ``DeepseekVLV2Config`` for parsing the nested config.
+
+from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
+
+
+class UnlimitedOCRConfig(DeepseekVLV2Config):
+    model_type = "unlimited-ocr"
+
+    # An explicit ``__init__`` is required: Transformers v5 processes each
+    # concrete config class' ``__init__`` signature to build nested sub-configs,
+    # and an empty subclass (only overriding ``model_type``) would skip
+    # ``DeepseekVLV2Config.__init__``, leaving ``text_config`` unset.
+    def __init__(
+        self,
+        tile_tag: str = "2D",
+        global_view_pos: str = "head",
+        candidate_resolutions: tuple[tuple[int, int]] = ((384, 384),),
+        rswa_window: int = 128,
+        **kwargs,
+    ):
+        super().__init__(
+            tile_tag=tile_tag,
+            global_view_pos=global_view_pos,
+            candidate_resolutions=candidate_resolutions,
+            **kwargs,
+        )
+        self.rswa_window = rswa_window

vllm/transformers_utils/model_arch_config_convertor.py

@@ -310,6 +310,12 @@ class ModelArchConfigConvertorBase:
             return False
         return self.hf_config.model_type in MM_PREFIX_LM_MODELS
 
+    def rswa_window(self) -> int | None:
+        value = getattr(self.hf_config, "rswa_window", None)
+        if value is None:
+            return None
+        return int(value)
+
     def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
         derived_max_model_len = float("inf")
         possible_keys = [
@@ -360,6 +366,7 @@ class ModelArchConfigConvertorBase:
             quantization_config=self.get_quantization_config(),
             is_deepseek_mla=self.is_deepseek_mla(),
             is_mm_prefix_lm=self.is_mm_prefix_lm(),
+            rswa_window=self.rswa_window(),
             derived_max_model_len_and_key=self.derive_max_model_len_and_key(),
         )
 

vllm/transformers_utils/processors/deepseek_ocr.py

@@ -161,10 +161,12 @@ class DeepseekOCRProcessor(ProcessorMixin):
         image_size: int = IMAGE_SIZE,
         base_size: int = BASE_SIZE,
         strategy: Literal["v1", "v2"] = "v1",
+        max_crops: int = MAX_CROPS,
         **kwargs,
     ):
         self.image_size = image_size
         self.base_size = base_size
+        self.max_crops = max_crops
 
         # image token calculation strategy for
         # Deepseek-OCR and Deepseek-OCR-2
@@ -332,7 +334,7 @@ class DeepseekOCRProcessor(ProcessorMixin):
                 crop_ratio = [1, 1]
             elif cropping:
                 images_crop_raw, crop_ratio = dynamic_preprocess(
-                    image, image_size=self.image_size
+                    image, image_size=self.image_size, max_num=self.max_crops
                 )
             else:
                 crop_ratio = [1, 1]

vllm/transformers_utils/processors/unlimited_ocr.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Image processor for Unlimited-OCR (baidu/Unlimited-OCR)."""
+
+from PIL import Image
+
+from vllm.logger import init_logger
+from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor
+
+logger = init_logger(__name__)
+
+
+class UnlimitedOCRProcessor(DeepseekOCRProcessor):
+    """DeepseekOCRProcessor variant for Unlimited-OCR.
+
+    The only behavioural difference from the base processor is a multi-image
+    safeguard: when more than one image is present, crop ("gundam") mode is
+    disabled.
+
+    Because the effective crop flag then depends on *how many* images are in the
+    request, the per-item processing output is no longer invariant of sibling
+    images. ``UnlimitedOCRMultiModalProcessor`` accounts for this by bypassing
+    the multimodal processing cache for multi-image requests (see its
+    ``_cached_apply_hf_processor``), so the two paths stay consistent.
+
+    DeepSeek-OCR does *not* have this restriction because its ``max_crops=6`` is
+    small enough to be safe for multi-image use.
+    """
+
+    def tokenize_with_images(
+        self,
+        conversation: str,
+        images: list[Image.Image],
+        bos: bool = True,
+        eos: bool = True,
+        cropping: bool = True,
+    ):
+        if len(images) > 1 and cropping:
+            logger.warning_once(
+                "Unlimited-OCR: crop mode is not supported for multi-image "
+                "input. Falling back to cropping=False."
+            )
+            cropping = False
+        return super().tokenize_with_images(
+            conversation, images, bos=bos, eos=eos, cropping=cropping
+        )

vllm/v1/attention/backend.py

@@ -455,6 +455,13 @@ class CommonAttentionMetadata:
     where bidirectional attention should apply. None for text-only
     batches or non-PrefixLM models."""
 
+    rswa_prefix_lens: torch.Tensor | None = None
+    """(batch_size,) per-request prefix length (prompt/image token count) for
+    Reference Sliding Window Attention (R-SWA). Tokens with logical index below
+    this stay globally visible; later (generated) tokens additionally see a
+    fixed sliding window. None disables R-SWA. The attention backend copies this
+    into its own persistent buffer and reads ``rswa_window`` from model config."""
+
     # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
     _seq_lens_cpu: torch.Tensor | None = None
     _num_computed_tokens_cpu: torch.Tensor | None = None
@@ -539,6 +546,7 @@ class CommonAttentionMetadata:
             dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
             dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
             is_prefilling=maybe_slice_reqs(self.is_prefilling),
+            rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),
         )
 
 

vllm/v1/attention/backends/flash_attn.py

@@ -256,6 +256,16 @@ class FlashAttentionMetadata:
     # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.
     mm_prefix_range_tensor: torch.Tensor | None = None
 
+    # Reference Sliding Window Attention (R-SWA) fields.
+    # rswa_prefix_lens:  per-request prompt lengths [num_reqs], int32, CUDA.
+    # rswa_window:       sliding window size (scalar int, for logic checks).
+    # rswa_window_tensor: [1] int32 CUDA tensor — pre-allocated in build() so
+    #   no CPU→CUDA copy is needed inside forward() during CUDA graph capture.
+    # Only populated when the model uses R-SWA (Unlimited-OCR).
+    rswa_prefix_lens: torch.Tensor | None = None
+    rswa_window: int | None = None
+    rswa_window_tensor: torch.Tensor | None = None
+
 
 def _get_sliding_window_configs(
     vllm_config: VllmConfig,
@@ -386,6 +396,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         # populated on first build() call.
         self.aot_sliding_window: tuple[int, int] | None = None
 
+        # R-SWA: persistent CUDA-graph-safe buffers owned by this builder.
+        self.rswa_window: int | None = self.model_config.rswa_window
+        self.persistent_rswa_prefix_lens: torch.Tensor | None = None
+        self.persistent_rswa_window_tensor: torch.Tensor | None = None
+        if self.rswa_window is not None:
+            max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+            self.persistent_rswa_prefix_lens = torch.zeros(
+                max_num_reqs, dtype=torch.int32, device=self.device
+            )
+            self.persistent_rswa_window_tensor = torch.tensor(
+                [self.rswa_window], dtype=torch.int32, device=self.device
+            )
+
     def build(
         self,
         common_prefix_len: int,
@@ -589,6 +612,22 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                 mm_ranges, num_reqs, seq_lens.device
             )
 
+        # R-SWA: copy prefix lengths into persistent buffers (outside the
+        # compiled region) so forward() never allocates during CUDA graph
+        # capture.  rswa_window is a static model config scalar read here.
+        if (
+            self.rswa_window is not None
+            and common_attn_metadata.rswa_prefix_lens is not None
+        ):
+            assert self.persistent_rswa_prefix_lens is not None
+            assert self.persistent_rswa_window_tensor is not None
+            src = common_attn_metadata.rswa_prefix_lens
+            rswa_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]
+            rswa_prefix_lens.copy_(src[:num_reqs], non_blocking=True)
+            attn_metadata.rswa_prefix_lens = rswa_prefix_lens
+            attn_metadata.rswa_window = self.rswa_window
+            attn_metadata.rswa_window_tensor = self.persistent_rswa_window_tensor
+
         return attn_metadata
 
     def update_block_table(
@@ -805,7 +844,7 @@ class FlashAttentionImpl(AttentionImpl):
                 )
                 return output
             else:
-                sliding_window_size = (
+                sliding_window_size: list[int] | None = (
                     list(self.sliding_window)
                     if self.sliding_window is not None
                     else None
@@ -840,6 +879,25 @@ class FlashAttentionImpl(AttentionImpl):
                     mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges)
                     mm_aux = [mm_prefix_ranges]
 
+                # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
+                # mask without block-size approximation.  The mask_mod encodes
+                # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which
+                # supersedes any FA-layer sliding_window_size parameter.
+                rswa_mask_mod_fn = None
+                rswa_aux = None
+                if (
+                    attn_metadata.rswa_prefix_lens is not None
+                    and self.vllm_flash_attn_version == 4
+                    and not is_dynamic_causal
+                ):
+                    rswa_mask_mod_fn = _make_rswa_mask_mod()
+                    rswa_aux = [
+                        attn_metadata.rswa_prefix_lens.to(torch.int32),
+                        attn_metadata.rswa_window_tensor,  # pre-allocated CUDA tensor
+                    ]
+                    # mask_mod fully expresses R-SWA; disable FA's own window.
+                    sliding_window_size = None
+
                 dynamic_causal = None
                 if isinstance(causal, torch.Tensor):
                     if self.vllm_flash_attn_version != 4:
@@ -873,8 +931,8 @@ class FlashAttentionImpl(AttentionImpl):
                     dynamic_causal=dynamic_causal,
                     num_splits=attn_metadata.max_num_splits,
                     s_aux=self.sinks,
-                    mask_mod=mm_mask_mod,
-                    aux_tensors=mm_aux,
+                    mask_mod=rswa_mask_mod_fn or mm_mask_mod,
+                    aux_tensors=rswa_aux or mm_aux,
                 )
                 return output
 
@@ -1152,6 +1210,59 @@ def _make_mm_prefix_mask_mod(max_ranges: int):
     return mm_prefix_mask_mod
 
 
+def _make_rswa_mask_mod():
+    """Build a CuTE-DSL mask_mod for Reference Sliding Window Attention (R-SWA).
+
+    FA4 varlen + paged-KV convention (verified from cute/mask.py apply_mask):
+      q_idx  = LOCAL query-token offset (0 .. seqlen_q - 1) within this sequence.
+      kv_idx = LOCAL KV-token position (0 .. seqlen_k - 1) within this sequence.
+
+    To recover the ABSOLUTE token position (needed for causal and the sliding
+    window distance), use the standard offset:
+      abs_q = q_idx + (seqlen_k - seqlen_q)
+
+    R-SWA keep condition:
+      abs_q >= kv_idx                    (causal: KV at or before the query)
+      AND (kv_idx < prefix_len           (global prefix is always visible)
+           OR  abs_q - kv_idx < window)  (generated tokens: sliding window)
+
+    aux_tensors[0]: prefix_lens [num_reqs] int32 — per-request prefill length.
+    aux_tensors[1]: rswa_window [1]        int32 — decode sliding window size.
+
+    use_fast_sampling=True lets FA4 skip fully-masked KV blocks (gap blocks)
+    without loading their data.
+    """
+    import cutlass.cute as cute
+    from cutlass import Int32  # type: ignore[attr-defined]
+
+    from vllm.vllm_flash_attn.cute.utils import (  # type: ignore[import-untyped]
+        scalar_to_ssa,
+    )
+
+    @cute.jit
+    def rswa_mask_mod(
+        batch_idx: cute.TensorSSA,
+        head_idx: cute.TensorSSA,
+        q_idx: cute.TensorSSA,
+        kv_idx: cute.TensorSSA,
+        seqlen_info,
+        aux_tensors,
+    ):
+        b = batch_idx[0]
+        prefix_len = scalar_to_ssa(aux_tensors[0][b], Int32)
+        window = scalar_to_ssa(aux_tensors[1][0], Int32)
+        # Convert local q offset to absolute token position.
+        offset = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
+        abs_q = q_idx + offset
+        causal = kv_idx <= abs_q
+        in_prefix = kv_idx < prefix_len
+        in_window = (abs_q - kv_idx) < window
+        return causal & (in_prefix | in_window)
+
+    rswa_mask_mod.use_fast_sampling = True
+    return rswa_mask_mod
+
+
 def use_cascade_attention(
     common_prefix_len: int,
     query_lens: np.ndarray,

vllm/v1/attention/backends/flex_attention.py

@@ -408,6 +408,11 @@ class FlexAttentionMetadata:
     sliding_window: int | None = None
     mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
     block_sparsity_hint: BlockSparsityHint | None = None
+    # Reference Sliding Window Attention (R-SWA): per-request prefix length
+    # (prompt/image tokens stay globally visible) plus a sliding window over
+    # generated tokens. Both must be set to enable.
+    rswa_prefix_lens: torch.Tensor | None = None
+    rswa_window: int | None = None
 
     @cached_property
     def logical_block_ids(self):
@@ -571,6 +576,52 @@ class FlexAttentionMetadata:
 
         return final_mask_mod
 
+    def get_rswa_mask_mod(self) -> _mask_mod_signature:
+        """Creates the Reference Sliding Window Attention (R-SWA) mask_mod.
+
+        R-SWA keeps the whole prefix (image + prompt tokens, i.e. logical index
+        ``< prefix_len``) globally visible while generated tokens additionally
+        attend a fixed sliding window of recent tokens. This term is combined
+        with the base causal mask via logical AND, so it only ever *removes*
+        far-away generated tokens that fall outside the window and outside the
+        prefix.
+        """
+
+        assert self.doc_ids is not None
+        assert self.rswa_prefix_lens is not None
+        assert self.rswa_window is not None
+        doc_ids = self.doc_ids
+        prefix_lens = self.rswa_prefix_lens
+        window = self.rswa_window
+
+        def rswa_mask_mod(
+            q_req: torch.Tensor,
+            logical_q_idx: torch.Tensor,
+            logical_kv_idx: torch.Tensor,
+        ) -> torch.Tensor:
+            prefix_len = prefix_lens[q_req]
+            in_prefix = logical_kv_idx < prefix_len
+            in_window = (logical_q_idx - logical_kv_idx) < window
+            return in_prefix | in_window
+
+        def final_mask_mod(
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ) -> torch.Tensor:
+            (is_valid, logical_q_idx, logical_kv_idx) = (
+                self._convert_physical_to_logical(doc_ids, q_idx, physical_kv_idx)
+            )
+            q_req = doc_ids[q_idx]
+            return torch.where(
+                is_valid,
+                rswa_mask_mod(q_req, logical_q_idx, logical_kv_idx),
+                False,
+            )
+
+        return final_mask_mod
+
     def get_mask_mod(self):
         # Stage-1: initialize the base mask_mod
         # (causal mask for decoder or bidirectional mask for encoder)
@@ -588,6 +639,10 @@ class FlexAttentionMetadata:
             # Add prefix LM mask for vision-language prefix LM attention
             prefix_lm_mask_mod = self.get_prefix_lm_mask_mod()
             mask_mod = or_masks(mask_mod, prefix_lm_mask_mod)
+        if self.rswa_window is not None and self.rswa_prefix_lens is not None:
+            # Reference Sliding Window Attention: AND with the base causal mask
+            # (prefix stays global, generated tokens use a sliding window).
+            mask_mod = and_masks(mask_mod, self.get_rswa_mask_mod())
         return mask_mod
 
     def get_transformed_score_mod(self) -> _score_mod_signature | None:
@@ -663,9 +718,21 @@ class FlexAttentionMetadata:
             self.doc_ids, : cdiv(self.max_seq_len, self.block_size)
         ]
 
-        custom_hint = self.block_sparsity_hint is not None
+        # block_table slots beyond each request's seq_len may contain garbage
+        # physical page ids (see physical_to_logical_mapping). With batched
+        # decode, max_seq_len is the batch max while shorter requests still
+        # index all columns up to that max unless masked here.
+        num_blocks = self.num_blocks_per_seq[self.doc_ids]
+        past_seq = self.logical_block_ids[None, :] >= num_blocks[:, None]
+        used_pages.masked_fill_(past_seq, 0)
 
-        if self.sliding_window or custom_hint:
+        custom_hint = self.block_sparsity_hint is not None
+        use_rswa = self.rswa_window is not None and self.rswa_prefix_lens is not None
+        needs_per_q_pruning = (
+            self.causal or self.sliding_window or custom_hint or use_rswa
+        )
+
+        if needs_per_q_pruning:
             device = used_pages.device
             assert self.doc_ids is not None
             token_indices = torch.arange(
@@ -676,6 +743,12 @@ class FlexAttentionMetadata:
                 - self.query_start_loc[self.doc_ids]
                 + self.decode_offset[self.doc_ids]
             )
+            block_starts = self.logical_block_ids * self.block_size
+            block_ends = block_starts + self.block_size
+
+            if self.causal:
+                future_blocks = block_starts[None, :] > logical_q_idx[:, None]
+                used_pages.masked_fill_(future_blocks, 0)
 
             if self.sliding_window:
                 assert self.sliding_window is not None
@@ -685,6 +758,23 @@ class FlexAttentionMetadata:
                 min_block_idx = min_kv_idx // self.block_size
                 sliding_mask = self.logical_block_ids >= min_block_idx[:, None]
                 used_pages.masked_fill_(~sliding_mask, 0)
+            if use_rswa:
+                # R-SWA keeps prefix KV globally visible and applies a sliding
+                # window over generated tokens. Prune blocks that fall entirely
+                # in the "hole" between prefix_len and the current window so
+                # FlexAttention does not gather invalid paged-KV slots (this
+                # mirrors uniform sliding-window block pruning above).
+                assert self.rswa_prefix_lens is not None
+                assert self.rswa_window is not None
+                prefix_len = self.rswa_prefix_lens[self.doc_ids]
+                min_kv_window = torch.maximum(
+                    prefix_len,
+                    logical_q_idx - (self.rswa_window - 1),
+                )
+                in_gap = (block_starts[None, :] >= prefix_len[:, None]) & (
+                    block_ends[None, :] <= min_kv_window[:, None]
+                )
+                used_pages.masked_fill_(in_gap, 0)
             if custom_hint:
                 assert self.block_sparsity_hint is not None
                 q_block_idx = logical_q_idx // self.block_size
@@ -798,12 +888,36 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
         self.max_num_query_groups = cdiv(max_num_batched_tokens, self.q_block_size)
         max_num_pages_per_seq = cdiv(self.max_model_len, self.block_size)
         self.max_num_kv_indices = self.q_block_size * max_num_pages_per_seq
+        # R-SWA uses q_block_size=1 so block lists are not merged across requests
+        # in a q-group (mixed-length batches otherwise gather foreign paged-KV).
+        self.max_num_rswa_query_groups = max_num_batched_tokens
+        # +1 sentinel column: the flex-attention kernel's get_offset_for_next_block
+        # always prefetches kv_indices[q, kv_num_blocks] (one past the last valid
+        # entry) to compute the jump offset for the next loop iteration.  When
+        # kv_num_blocks[q] == W (every page of the sequence is live), that prefetch
+        # reads column W of the persistent buffer.  Without the extra column this
+        # would land on stale data from a previous step (the buffer is wider than W
+        # but is never fully zeroed), producing an out-of-bounds K/V pointer and a
+        # CUDA illegal memory access.  Allocating W_max+1 columns and initialising
+        # the whole buffer to -1 ensures the sentinel slot is always safe to read.
+        self.max_num_rswa_kv_indices = max_num_pages_per_seq + 1
         self.persistent_kv_num_blocks = torch.empty(
             self.max_num_query_groups, dtype=torch.int32, device=device
         )
+        self.persistent_rswa_kv_num_blocks = torch.empty(
+            self.max_num_rswa_query_groups, dtype=torch.int32, device=device
+        )
         self.persistent_offset_tensor = torch.empty(
             max_num_seqs, dtype=torch.int32, device=device
         )
+        # Persistent buffer for R-SWA per-request prefix lengths so the device
+        # address stays stable across steps (required for CUDA graph replay).
+        self.rswa_window: int | None = self.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(
+                max_num_seqs, dtype=torch.int32, device=device
+            )
         self.persistent_doc_ids = torch.empty(
             max_num_batched_tokens, dtype=torch.int32, device=device
         )
@@ -811,6 +925,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
         # initialize later when we can access block_table
         self.persistent_physical_to_logical = None
         self.persistent_kv_indices = None
+        self.persistent_rswa_kv_indices = None
 
         self.custom_logical_mask_mod: _mask_mod_signature | None = None
         if self._uses_full_cudagraphs():
@@ -936,6 +1051,26 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
                 dtype=torch.int32,
                 device=self.device,
             )
+        if self.persistent_rswa_kv_indices is None:
+            # Initialise to -1 so the +1 sentinel column (see max_num_rswa_kv_indices)
+            # is always a safe pad value for the flex kernel's prefetch.
+            self.persistent_rswa_kv_indices = torch.full(
+                (self.max_num_rswa_query_groups, self.max_num_rswa_kv_indices),
+                fill_value=-1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+
+        use_rswa = self.rswa_window is not None
+        q_block_size = 1 if use_rswa else self.q_block_size
+        persistent_kv_indices = (
+            self.persistent_rswa_kv_indices if use_rswa else self.persistent_kv_indices
+        )
+        persistent_kv_num_blocks = (
+            self.persistent_rswa_kv_num_blocks
+            if use_rswa
+            else self.persistent_kv_num_blocks
+        )
 
         inverse_block_table = copy_to_persistent(
             self.persistent_physical_to_logical, inverse_block_table
@@ -944,6 +1079,13 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
         offset_tensor = common_attn_metadata.compute_num_computed_tokens()
         offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor)
 
+        rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens
+        if use_rswa and rswa_prefix_lens is not None:
+            assert self.persistent_rswa_prefix_lens is not None
+            rswa_prefix_lens = copy_to_persistent(
+                self.persistent_rswa_prefix_lens, rswa_prefix_lens
+            )
+
         uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec)
         logical_mask_mod = (
             bidirectional_mask_mod
@@ -986,12 +1128,14 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
             # attention block mask for encoder-only models, disable it temporarily.
             # see: https://github.com/vllm-project/vllm/pull/27329#issuecomment-3431484053
             direct_build=self.direct_build and uses_paged_kv,
-            q_block_size=self.q_block_size,
+            q_block_size=q_block_size,
             kv_block_size=self.kv_block_size,
-            persistent_kv_indices=self.persistent_kv_indices,
-            persistent_kv_num_blocks=self.persistent_kv_num_blocks,
+            persistent_kv_indices=persistent_kv_indices,
+            persistent_kv_num_blocks=persistent_kv_num_blocks,
             persistent_doc_ids=self.persistent_doc_ids,
             mm_prefix_range=common_attn_metadata.mm_req_doc_ranges,
+            rswa_prefix_lens=rswa_prefix_lens,
+            rswa_window=self.rswa_window,
         )
 
         # Pre-build block_mask so it is ready before CUDA graph capture.

vllm/v1/core/kv_cache_coordinator.py

@@ -329,7 +329,10 @@ class KVCacheCoordinator(ABC):
         ]
 
     def remove_skipped_blocks(
-        self, request_id: str, total_computed_tokens: int
+        self,
+        request_id: str,
+        total_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
     ) -> None:
         """
         Remove the blocks that are no longer needed from `blocks` and replace
@@ -339,9 +342,14 @@ class KVCacheCoordinator(ABC):
             request_id: The request ID.
             total_computed_tokens: The total number of computed tokens, including
                 local computed tokens and external computed tokens.
+            num_prompt_tokens: Optional prompt length. R-SWA managers use this to
+                free gap blocks between the prefill tail and decode window; other
+                manager types ignore it.
         """
         for manager in self.single_type_managers:
-            manager.remove_skipped_blocks(request_id, total_computed_tokens)
+            manager.remove_skipped_blocks(
+                request_id, total_computed_tokens, num_prompt_tokens
+            )
 
     def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
         """

vllm/v1/core/kv_cache_manager.py

@@ -398,7 +398,9 @@ class KVCacheManager:
         # Should call this function before allocating new blocks to reduce
         # the number of evicted blocks.
         self.coordinator.remove_skipped_blocks(
-            request.request_id, total_computed_tokens
+            request.request_id,
+            total_computed_tokens,
+            num_prompt_tokens=request.num_prompt_tokens,
         )
 
         num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
@@ -468,7 +470,10 @@ class KVCacheManager:
         self.coordinator.free(request.request_id)
 
     def remove_skipped_blocks(
-        self, request_id: str, total_computed_tokens: int
+        self,
+        request_id: str,
+        total_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
     ) -> None:
         """Remove the blocks that are no longer needed from `blocks` and replace
         the removed blocks with null_block.
@@ -477,8 +482,11 @@ class KVCacheManager:
             request_id: The request ID.
             total_computed_tokens: The total number of computed tokens, including
                 local computed tokens and external computed tokens.
+            num_prompt_tokens: Optional prompt length for R-SWA gap eviction.
         """
-        self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens)
+        self.coordinator.remove_skipped_blocks(
+            request_id, total_computed_tokens, num_prompt_tokens
+        )
 
     def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
         """Pop the request's bookkeeping and return its blocks without

vllm/v1/core/sched/scheduler.py

@@ -2342,6 +2342,7 @@ class Scheduler(SchedulerInterface):
         self.kv_cache_manager.remove_skipped_blocks(
             request_id=request.request_id,
             total_computed_tokens=request.num_computed_tokens,
+            num_prompt_tokens=request.num_prompt_tokens,
         )
 
         block_ids = self.kv_cache_manager.get_block_ids(request.request_id)

vllm/v1/core/single_type_kv_cache_manager.py

@@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheSpec,
     MambaSpec,
     MLAAttentionSpec,
+    RSWASpec,
     SinkFullAttentionSpec,
     SlidingWindowMLASpec,
     SlidingWindowSpec,
@@ -476,8 +477,38 @@ class SingleTypeKVCacheManager(ABC):
 
         raise NotImplementedError
 
+    def _remove_blocks_in_range(
+        self,
+        request_id: str,
+        first_block: int,
+        last_block: int,
+    ) -> None:
+        """Free blocks in ``[first_block, last_block)`` and replace with null_block.
+
+        Iterates backward so newly-evictable tail blocks are reached even after
+        earlier blocks in the range were nulled in a prior call.
+        """
+        if request_id not in self.req_to_blocks:
+            return
+        if first_block >= last_block:
+            return
+        blocks = self.req_to_blocks[request_id]
+        last_block = min(last_block, len(blocks))
+
+        freed: list[KVCacheBlock] = []
+        for i in range(last_block - 1, first_block - 1, -1):
+            if blocks[i] == self._null_block:
+                break
+            freed.append(blocks[i])
+            blocks[i] = self._null_block
+        if freed:
+            self.block_pool.free_blocks(freed)
+
     def remove_skipped_blocks(
-        self, request_id: str, total_computed_tokens: int
+        self,
+        request_id: str,
+        total_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
     ) -> None:
         """
         Remove and free the blocks that are no longer needed for attention computation.
@@ -490,7 +521,11 @@ class SingleTypeKVCacheManager(ABC):
             request_id: The request ID.
             total_computed_tokens: The total number of computed tokens, including
                 local computed tokens and external computed tokens.
+            num_prompt_tokens: Optional prompt length for attention types (e.g.
+                R-SWA) that evict a middle gap rather than a head prefix. Ignored
+                by the default implementation.
         """
+        del num_prompt_tokens
         # Remove the blocks that will be skipped during attention computation.
         num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
         if num_skipped_tokens <= 0:
@@ -506,18 +541,7 @@ class SingleTypeKVCacheManager(ABC):
         # range), so we must cap to the number of blocks that currently exist for
         # this request.
         num_skipped_blocks = min(num_skipped_blocks, len(blocks))
-        removed_blocks: list[KVCacheBlock] = []
-        # Because the block starts from index 0, the num_skipped_block-th block
-        # corresponds to index num_skipped_blocks - 1.
-        for i in range(num_skipped_blocks - 1, -1, -1):
-            if blocks[i] == self._null_block:
-                # If the block is already a null block, the blocks before it
-                # should also have been set to null blocks by the previous calls
-                # to this function.
-                break
-            removed_blocks.append(blocks[i])
-            blocks[i] = self._null_block
-        self.block_pool.free_blocks(removed_blocks)
+        self._remove_blocks_in_range(request_id, 0, num_skipped_blocks)
 
     def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
         """
@@ -598,6 +622,50 @@ class FullAttentionManager(SingleTypeKVCacheManager):
         return num_common_blocks
 
 
+class RSWAManager(FullAttentionManager):
+    """KV cache manager for Reference Sliding Window Attention (R-SWA).
+
+    When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees
+    gap blocks between the prefill tail and the current decode window.  This
+    bounds per-request KV memory at O(prefix_len + rswa_window) instead of
+    growing linearly with decode length.
+    """
+
+    def __init__(self, kv_cache_spec: RSWASpec, **kwargs) -> None:
+        super().__init__(kv_cache_spec, **kwargs)
+        self.rswa_window: int = kv_cache_spec.rswa_window
+
+    def remove_skipped_blocks(
+        self,
+        request_id: str,
+        total_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
+    ) -> None:
+        """Free gap blocks that are no longer needed for attention.
+
+        Gap = blocks entirely within
+            [ceil(prefix_len / block_size) * block_size,
+             max(prefix_len, total_computed_tokens - rswa_window))
+
+        Freed blocks are replaced with null_block in req_to_blocks so the
+        block_table passed to FA4 is valid (null_block KV is all-zero;
+        rswa_mask_mod marks gap positions as non-visible so FA4 skips them).
+        """
+        if num_prompt_tokens is None:
+            super().remove_skipped_blocks(
+                request_id, total_computed_tokens, num_prompt_tokens
+            )
+            return
+
+        bs = self.block_size
+        # First block fully after the prefill boundary.
+        first_gap_block = cdiv(num_prompt_tokens, bs)
+        # Decode window start position; blocks before this are evictable.
+        window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window)
+        last_gap_block = window_start // bs  # exclusive upper bound
+        self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block)
+
+
 class SlidingWindowManager(SingleTypeKVCacheManager):
     def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:
         super().__init__(kv_cache_spec, **kwargs)
@@ -1072,7 +1140,12 @@ class MambaManager(SingleTypeKVCacheManager):
 
         return mask
 
-    def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
+    def remove_skipped_blocks(
+        self,
+        request_id: str,
+        num_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
+    ) -> None:
         assert isinstance(self.kv_cache_spec, MambaSpec)
 
         # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain
@@ -1082,7 +1155,9 @@ class MambaManager(SingleTypeKVCacheManager):
         # that we might actually need.
         num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks)
 
-        super().remove_skipped_blocks(request_id, num_computed_tokens)
+        super().remove_skipped_blocks(
+            request_id, num_computed_tokens, num_prompt_tokens
+        )
         if self.mamba_cache_mode == "align":
             # `last_state_block_idx` refers to the block index allocated two steps ago.
             # The block allocated in the previous step is used to copy Mamba states
@@ -1401,10 +1476,16 @@ def get_manager_for_kv_cache_spec(
     assert manager_class is not None, (
         f"No manager registered for KVCacheSpec {type(kv_cache_spec)}"
     )
-    # SlidingWindow / ChunkedLocalAttention managers recycle blocks across
-    # chunks; the runtime admission cap must match the recycling-aware bound
-    # the startup pool sizer uses (single source of truth: the spec method).
-    if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):
+    # SlidingWindow / ChunkedLocalAttention managers recycle blocks;
+    # the runtime admission cap must match the recycling-aware bound the
+    # startup pool sizer uses (single source of truth: the spec method).
+    # R-SWA also recycles gap blocks but peak physical KV still fits the
+    # full-attention bound (prefix + window <= max_model_len), so it inherits
+    # FullAttentionSpec sizing without a separate admission cap.
+    if isinstance(
+        kv_cache_spec,
+        (SlidingWindowSpec, ChunkedLocalAttentionSpec),
+    ):
         kwargs["max_admission_blocks_per_request"] = (
             kv_cache_spec.max_admission_blocks_per_request(
                 max_num_batched_tokens=max_num_batched_tokens,
@@ -1457,6 +1538,9 @@ def register_all_kvcache_specs(vllm_config):
     KVCacheSpecRegistry.register(
         MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec
     )
+    KVCacheSpecRegistry.register(
+        RSWASpec, RSWAManager, uniform_type_base_spec=FullAttentionSpec
+    )
     # NOTE(Mengqing): HiddenStateCacheSpec won't take part in
     # grouping, thus the uniform_type_base_spec is just a
     # placeholder.

vllm/v1/kv_cache_interface.py

@@ -437,6 +437,46 @@ class HiddenStateCacheSpec(MLAAttentionSpec):
     pass
 
 
+@dataclass(frozen=True, kw_only=True)
+class RSWASpec(FullAttentionSpec):
+    """KV cache spec for Reference Sliding Window Attention (R-SWA).
+
+    Prefill (image + text prompt) tokens are always globally visible.
+    Only the last ``rswa_window`` generated tokens are kept in the KV cache;
+    gap blocks (between the prefill tail and the current decode window) are
+    evicted during each decode step to bound memory at
+    O(prefix_blocks + window_blocks).
+    """
+
+    rswa_window: int
+
+    @classmethod
+    def merge(cls, specs: list[RSWASpec]) -> RSWASpec:
+        assert all(isinstance(spec, RSWASpec) for spec in specs), (
+            "All attention layers in the same KV cache group must be RSWASpec."
+        )
+        rswa_windows = {spec.rswa_window for spec in specs}
+        assert len(rswa_windows) == 1, (
+            f"All R-SWA layers must share the same rswa_window, got {rswa_windows}"
+        )
+        # Delegate common field merging to the parent, then reattach rswa_window.
+        base = FullAttentionSpec.merge(specs)  # type: ignore[arg-type]
+        return cls(
+            block_size=base.block_size,
+            num_kv_heads=base.num_kv_heads,
+            head_size=base.head_size,
+            head_size_v=base.head_size_v,
+            dtype=base.dtype,
+            kv_quant_mode=base.kv_quant_mode,
+            page_size_padded=base.page_size_padded,
+            indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,
+            sliding_window=base.sliding_window,
+            attention_chunk_size=base.attention_chunk_size,
+            non_causal=base.non_causal,
+            rswa_window=rswa_windows.pop(),
+        )
+
+
 @dataclass(frozen=True, kw_only=True)
 class ChunkedLocalAttentionSpec(AttentionSpec):
     attention_chunk_size: int

vllm/v1/worker/gpu/attn_utils.py

@@ -469,6 +469,7 @@ def build_attn_metadata(
     model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     for_cudagraph_capture: bool = False,
     causal: bool = True,
+    rswa_prefix_lens: torch.Tensor | None = None,
 ) -> dict[str, Any]:
     seq_lens = seq_lens[:num_reqs]
     if dcp_local_seq_lens is not None:
@@ -501,6 +502,7 @@ def build_attn_metadata(
             causal=causal,
             dcp_local_seq_lens=dcp_local_seq_lens,
             positions=positions,
+            rswa_prefix_lens=rswa_prefix_lens,
             **common_attn_metadata_extra_kwargs,
         )
 

vllm/v1/worker/gpu/input_batch.py

@@ -96,6 +96,9 @@ class InputBatch:
     # Whether any requests in batch use structured output.
     has_structured_output_reqs: bool
 
+    # [num_reqs_after_padding] per-request prompt length for R-SWA (optional).
+    rswa_prefix_lens: torch.Tensor | None = None
+
     @classmethod
     def make_dummy(
         cls,

vllm/v1/worker/gpu/model_runner.py

@@ -223,6 +223,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             max_num_tokens=self.max_num_tokens,
             device=self.device,
         )
+        # R-SWA: persistent GPU buffer for per-request prefix lengths (CUDA-graph safe).
+        self.rswa_prefix_lens_buffer: torch.Tensor | None = None
+        if self.model_config.rswa_window is not None:
+            self.rswa_prefix_lens_buffer = torch.zeros(
+                self.max_num_reqs, dtype=torch.int32, device=self.device
+            )
 
         if self.use_pp:
             self.pp_handler = PPHandler(
@@ -985,6 +991,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         if self.use_pp:
             # max_seq_len is only consumed by the PP `compute_need_sampled_mask`
             max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
+
+        rswa_prefix_lens = None
+        if self.rswa_prefix_lens_buffer is not None:
+            rswa_prefix_lens = self.rswa_prefix_lens_buffer[:num_reqs_padded]
+            rswa_prefix_lens[:num_reqs] = self.req_states.prompt_len.gpu[
+                idx_mapping[:num_reqs]
+            ]
+            if num_reqs_padded > num_reqs:
+                rswa_prefix_lens[num_reqs:].zero_()
+
         return InputBatch(
             req_ids=req_ids,
             num_reqs=num_reqs,
@@ -1015,6 +1031,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
+            rswa_prefix_lens=rswa_prefix_lens,
         )
 
     def prepare_attn(

vllm/v1/worker/gpu/model_states/default.py

@@ -168,5 +168,6 @@ class DefaultModelState(ModelState):
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             positions=input_batch.positions,
             for_cudagraph_capture=for_capture,
+            rswa_prefix_lens=input_batch.rswa_prefix_lens,
         )
         return attn_metadata

vllm/v1/worker/gpu/model_states/encoder_decoder.py

@@ -146,6 +146,7 @@ class EncoderDecoderModelState(ModelState):
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             model_specific_attn_metadata=enc_dec_attn_metadata,
             for_cudagraph_capture=for_capture,
+            rswa_prefix_lens=input_batch.rswa_prefix_lens,
         )
         return attn_metadata
 

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -141,6 +141,7 @@ class MambaHybridModelState(DefaultModelState):
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             model_specific_attn_metadata=mamba_attn_metadata,
             for_cudagraph_capture=for_capture,
+            rswa_prefix_lens=input_batch.rswa_prefix_lens,
         )
 
     def postprocess_state(

vllm/v1/worker/gpu_model_runner.py

@@ -2335,6 +2335,13 @@ class GPUModelRunner(
                 req_idx = self.input_batch.req_id_to_index[req_id]
                 req_doc_ranges[req_idx] = image_doc_ranges
 
+        # Reference Sliding Window Attention (R-SWA): pass per-request prompt
+        # lengths so the attention backend can keep the prefix globally visible.
+        # The backend owns the persistent CUDA-graph-safe GPU buffer.
+        rswa_prefix_lens = None
+        if self.model_config.rswa_window is not None:
+            rswa_prefix_lens = num_prompt_tokens_cpu
+
         cm_base = CommonAttentionMetadata(
             query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],
             query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],
@@ -2352,6 +2359,7 @@ class GPUModelRunner(
             is_prefilling=is_prefilling,
             positions=self.positions[:num_tokens_padded],
             mm_req_doc_ranges=req_doc_ranges,
+            rswa_prefix_lens=rswa_prefix_lens,
         )
 
         if self.dcp_world_size > 1: