diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index b0e0e3ce9c4d..70b81aed9cd1 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -629,6 +629,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
| `Tarsier2ForConditionalGeneration`^ | Tarsier2 | T + IE+ + VE+ | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
| `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
+| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I+ | `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!
diff --git a/tests/models/registry.py b/tests/models/registry.py
index be271ea07771..463ce44851b5 100644
--- a/tests/models/registry.py
+++ b/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
),
diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py
index 7e960c2a6a34..609c1428d196 100644
--- a/tests/v1/core/test_single_type_kv_cache_manager.py
+++ b/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(
diff --git a/vllm/config/model.py b/vllm/config/model.py
index fecb26aa7e09..ef0600af54d5 100644
--- a/vllm/config/model.py
+++ b/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
diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py
index 0b99df22b880..0b4744de4898 100644
--- a/vllm/config/model_arch.py
+++ b/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."""
diff --git a/vllm/model_executor/layers/attention/__init__.py b/vllm/model_executor/layers/attention/__init__.py
index ca3574164d50..c9e477fb114f 100644
--- a/vllm/model_executor/layers/attention/__init__.py
+++ b/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",
]
diff --git a/vllm/model_executor/layers/attention/rswa_attention.py b/vllm/model_executor/layers/attention/rswa_attention.py
new file mode 100644
index 000000000000..c982722ff8ef
--- /dev/null
+++ b/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,
+ )
diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py
index ac6761498689..5c2278deb77e 100644
--- a/vllm/model_executor/models/config.py
+++ b/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,
}
diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py
index 0e061d6c6b5a..b811afafb0e8 100644
--- a/vllm/model_executor/models/deepseek_ocr.py
+++ b/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
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index 9b08ca9825e3..09960050c061 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/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): ...
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index 1f9e3a24fe48..dfc034729d8a 100644
--- a/vllm/model_executor/models/registry.py
+++ b/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",
diff --git a/vllm/model_executor/models/unlimited_ocr.py b/vllm/model_executor/models/unlimited_ocr.py
new file mode 100644
index 000000000000..06dc02512f10
--- /dev/null
+++ b/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)
diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py
index eb7f8b0cf0df..f90e427aee0d 100644
--- a/vllm/tokenizers/registry.py
+++ b/vllm/tokenizers/registry.py
@@ -31,7 +31,11 @@
# 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"),
diff --git a/vllm/transformers_utils/chat_templates/registry.py b/vllm/transformers_utils/chat_templates/registry.py
index a5f9bdac200a..ed744742903e 100644
--- a/vllm/transformers_utils/chat_templates/registry.py
+++ b/vllm/transformers_utils/chat_templates/registry.py
@@ -29,6 +29,7 @@ def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Path |
"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,
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index d63664072475..654d11df30d3 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -124,6 +124,7 @@ def __getitem__(self, key):
laguna="LagunaConfig",
lfm2_moe="Lfm2MoeConfig",
tarsier2="Tarsier2Config",
+ **{"unlimited-ocr": "UnlimitedOCRConfig"},
)
_SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"}
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index 021eb2ea419e..871cb5249005 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -73,6 +73,7 @@
"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 @@
"RadioConfig",
"SpeculatorsConfig",
"UltravoxConfig",
+ "UnlimitedOCRConfig",
"Step3VLConfig",
"Step3VisionEncoderConfig",
"Step3TextConfig",
diff --git a/vllm/transformers_utils/configs/unlimited_ocr.py b/vllm/transformers_utils/configs/unlimited_ocr.py
new file mode 100644
index 000000000000..99e50a03c7db
--- /dev/null
+++ b/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
diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py
index 483f1d8be81b..e372834d68d7 100644
--- a/vllm/transformers_utils/model_arch_config_convertor.py
+++ b/vllm/transformers_utils/model_arch_config_convertor.py
@@ -310,6 +310,12 @@ def is_mm_prefix_lm(self) -> bool:
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 @@ def convert(self) -> ModelArchitectureConfig:
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(),
)
diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py
index 68a2b1aaaa02..618070b506f7 100644
--- a/vllm/transformers_utils/processors/deepseek_ocr.py
+++ b/vllm/transformers_utils/processors/deepseek_ocr.py
@@ -161,10 +161,12 @@ def __init__(
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 @@ def tokenize_with_images(
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]
diff --git a/vllm/transformers_utils/processors/unlimited_ocr.py b/vllm/transformers_utils/processors/unlimited_ocr.py
new file mode 100644
index 000000000000..927f19d0f930
--- /dev/null
+++ b/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
+ )
diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py
index ccd70c6ca3ce..61a4e521c40d 100644
--- a/vllm/v1/attention/backend.py
+++ b/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 @@ def unpadded(
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),
)
diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py
index 75231bafeed5..c167708ac9c9 100755
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/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 @@ def __init__(
# 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 @@ def schedule(
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 @@ def forward(
)
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 @@ def forward(
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 @@ def forward(
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 mm_prefix_mask_mod(
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,
diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py
index 983544b56020..c45294bfc796 100644
--- a/vllm/v1/attention/backends/flex_attention.py
+++ b/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 @@ def final_mask_mod(
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 @@ def get_mask_mod(self):
# 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 @@ def _build_block_mask_direct(self) -> BlockMask:
self.doc_ids, : cdiv(self.max_seq_len, self.block_size)
]
+ # 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)
+
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 self.sliding_window or custom_hint:
+ if needs_per_q_pruning:
device = used_pages.device
assert self.doc_ids is not None
token_indices = torch.arange(
@@ -676,6 +743,12 @@ def _build_block_mask_direct(self) -> BlockMask:
- 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 @@ def _build_block_mask_direct(self) -> BlockMask:
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 @@ def __init__(
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 @@ def __init__(
# 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 @@ def build(
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 @@ def build(
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 @@ def build(
# 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.
diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py
index 48f597e1f249..a759d7a80add 100644
--- a/vllm/v1/core/kv_cache_coordinator.py
+++ b/vllm/v1/core/kv_cache_coordinator.py
@@ -329,7 +329,10 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
]
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 @@ def remove_skipped_blocks(
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], ...]:
"""
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index b0f6655bf957..57cd1490e81a 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -398,7 +398,9 @@ def allocate_slots(
# 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 @@ def free(self, request: Request) -> None:
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 @@ def remove_skipped_blocks(
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
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index ab9fd5e34331..ec479f093047 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2342,6 +2342,7 @@ def _connector_finished(
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)
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index e21c20a22819..642fe3e6a086 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -20,6 +20,7 @@
KVCacheSpec,
MambaSpec,
MLAAttentionSpec,
+ RSWASpec,
SinkFullAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
@@ -476,8 +477,38 @@ def find_longest_cache_hit(
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 @@ def remove_skipped_blocks(
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 @@ def remove_skipped_blocks(
# 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 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
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 @@ def reachable_block_mask(
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 @@ def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> No
# 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.
diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py
index b312a0fbeef1..323b1e763a52 100644
--- a/vllm/v1/kv_cache_interface.py
+++ b/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
diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py
index 737feb7d2770..758bd3bac7a6 100644
--- a/vllm/v1/worker/gpu/attn_utils.py
+++ b/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,
)
diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
index d745dc6abf9d..a6a2b296e388 100644
--- a/vllm/v1/worker/gpu/input_batch.py
+++ b/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,
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 927ece4fbac5..ce1bb7f5504d 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -223,6 +223,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
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 @@ def prepare_inputs(
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 @@ def prepare_inputs(
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(
diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py
index 2e14eb2e7d96..f760fc36dea6 100644
--- a/vllm/v1/worker/gpu/model_states/default.py
+++ b/vllm/v1/worker/gpu/model_states/default.py
@@ -168,5 +168,6 @@ def prepare_attn(
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
diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
index 889e624623d3..9edda27538e0 100644
--- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py
+++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
@@ -146,6 +146,7 @@ def prepare_attn(
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
diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
index 329f008a4e30..e08b09f1895e 100644
--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
@@ -141,6 +141,7 @@ def prepare_attn(
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(
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 74938a823d9f..6af531157752 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -2335,6 +2335,13 @@ def _get_block_table(kv_cache_gid: int):
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 @@ def _get_block_table(kv_cache_gid: int):
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: