@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c6741b2ad4

[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>

Tianyu Guo committed Jun 28, 2026 at 14:09 UTC c6741b2ad48a46e87d2cce35d113c4ae0950af91

35 files changed +1083 -48

docs/models/supported_models.md

@@ -629,6 +629,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

629 | `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |

631 | `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |

632 +| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I+ | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ |

633

634 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!

635

tests/models/registry.py

@@ -825,6 +825,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

825 "DeepseekOCR2ForCausalLM": _HfExamplesInfo(

826 "deepseek-ai/DeepSeek-OCR-2",

827 ),

828 + "UnlimitedOCRForCausalLM": _HfExamplesInfo(

829 + "baidu/Unlimited-OCR",

830 + ),

831 "DotsOCRForCausalLM": _HfExamplesInfo(

832 "rednote-hilab/dots.ocr", trust_remote_code=True

833 ),

tests/v1/core/test_single_type_kv_cache_manager.py

+51 -1

@@ -14,9 +14,14 @@ from vllm.v1.core.kv_cache_utils import (

14 )

15 from vllm.v1.core.single_type_kv_cache_manager import (

16 ChunkedLocalAttentionManager,

17 + RSWAManager,

18 SlidingWindowManager,

19 )

19 -from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec

20 +from vllm.v1.kv_cache_interface import (

21 + ChunkedLocalAttentionSpec,

22 + RSWASpec,

23 + SlidingWindowSpec,

24 +)

26 pytestmark = pytest.mark.cpu_test

@@ -327,6 +332,51 @@ def test_sliding_window_remove_skipped_blocks():

332 assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])

333

334

335 +def test_rswa_remove_skipped_blocks_gap_range():

336 + block_size = 4

337 + rswa_spec = RSWASpec(

338 + block_size=block_size,

339 + num_kv_heads=1,

340 + head_size=1,

341 + dtype=torch.float32,

342 + rswa_window=8,

343 + )

344 + block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=4)

345 + manager = RSWAManager(

346 + rswa_spec,

347 + block_pool=block_pool,

348 + enable_caching=True,

349 + kv_cache_group_id=0,

350 + scheduler_block_size=block_size,

351 + )

352 +

353 + null_block_id = block_pool.null_block.block_id

354 + original_block_ids = list(range(1000, 1010))

355 + block_table = [

356 + KVCacheBlock(id_) if id_ != null_block_id else block_pool.null_block

357 + for id_ in original_block_ids

358 + ]

359 + manager.req_to_blocks["test"] = block_table

360 +

361 + prefix_len = 16

362 +

363 + # Without num_prompt_tokens, R-SWA does not evict gap blocks.

364 + manager.remove_skipped_blocks("test", 28)

365 + assert [b.block_id for b in block_table] == original_block_ids

366 +

367 + # Gap = block 4 only (tokens [16, 20) fall in the gap).

368 + manager.remove_skipped_blocks("test", 28, num_prompt_tokens=prefix_len)

369 + expected = original_block_ids.copy()

370 + expected[4] = null_block_id

371 + assert [b.block_id for b in block_table] == expected

372 +

373 + # Window moves: blocks 5 and 6 also enter the gap; block 4 is already null.

374 + manager.remove_skipped_blocks("test", 36, num_prompt_tokens=prefix_len)

375 + expected[5] = null_block_id

376 + expected[6] = null_block_id

377 + assert [b.block_id for b in block_table] == expected

378 +

379 +

380 def test_get_num_blocks_to_allocate():

381 block_size = 2

382 sliding_window_spec = SlidingWindowSpec(

vllm/config/model.py

@@ -1251,6 +1251,10 @@ class ModelConfig:

1251 def is_mm_prefix_lm(self) -> bool:

1252 return self.model_arch_config.is_mm_prefix_lm

1253

1254 + @property

1255 + def rswa_window(self) -> int | None:

1256 + return self.model_arch_config.rswa_window

1257 +

1258 def get_head_size(self) -> int:

1259 return self.model_arch_config.head_size

1260

vllm/config/model_arch.py

@@ -56,5 +56,8 @@ class ModelArchitectureConfig:

56 is_mm_prefix_lm: bool

57 """Whether the model uses image bidirectional attention."""

59 + rswa_window: int | None

60 + """Reference Sliding Window Attention window size (None disables R-SWA)."""

61 +

62 derived_max_model_len_and_key: tuple[float, str | None]

63 """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

14 from vllm.model_executor.layers.attention.prefill_prefix_lm_attention import (

15 PrefillPrefixLMAttention,

16 )

17 +from vllm.model_executor.layers.attention.rswa_attention import RSWAAttention

18 from vllm.model_executor.layers.attention.static_sink_attention import (

19 StaticSinkAttention,

20 )

@@ -26,5 +27,6 @@ __all__ = [

27 "MLAAttention",

28 "MMEncoderAttention",

29 "PrefillPrefixLMAttention",

30 + "RSWAAttention",

31 "StaticSinkAttention",

32 ]

vllm/model_executor/layers/attention/rswa_attention.py new

+37

@@ -0,0 +1,37 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from vllm.config.vllm import VllmConfig

5 +from vllm.model_executor.layers.attention import Attention

6 +from vllm.v1.kv_cache_interface import KVCacheSpec, RSWASpec, get_kv_quant_mode

7 +

8 +

9 +class RSWAAttention(Attention):

10 + """Attention layer that reports ``RSWASpec`` as its KV cache spec.

11 +

12 + Drop-in replacement for the standard ``Attention`` layer when the model is

13 + configured with Reference Sliding Window Attention (R-SWA,

14 + ``rswa_window > 0``). The actual masking logic lives in the attention

15 + backend (FlexAttention or FA4 mask_mod); this layer only overrides

16 + ``get_kv_cache_spec`` so the KV cache manager instantiates ``RSWAManager``

17 + (instead of ``FullAttentionManager``) and can therefore evict "gap" blocks

18 + to keep per-request KV memory bounded at O(prefix + window).

19 + """

20 +

21 + def __init__(self, *args, rswa_window: int, **kwargs) -> None:

22 + super().__init__(*args, **kwargs)

23 + self._rswa_window = rswa_window

24 +

25 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:

26 + spec = super().get_kv_cache_spec(vllm_config)

27 + if spec is None:

28 + return None

29 + return RSWASpec(

30 + block_size=vllm_config.cache_config.block_size,

31 + num_kv_heads=self.num_kv_heads,

32 + head_size=self.head_size,

33 + head_size_v=self.head_size_v,

34 + dtype=self.kv_cache_torch_dtype,

35 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

36 + rswa_window=self._rswa_window,

37 + )

vllm/model_executor/models/config.py

+134

@@ -54,6 +54,139 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):

54 hf_config.is_causal = not hf_config.use_bidirectional_attention

57 +class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):

58 + @staticmethod

59 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:

60 + """Configure Unlimited-OCR attention backends for R-SWA and vision.

61 +

62 + Backend selection — controlled by the standard ``--attention-config``

63 + CLI argument (priority order):

64 +

65 + 1. ``--attention-config '{"backend": "FLASH_ATTN"}'``

66 + → FA4 + rswa_mask_mod. Exact token-level R-SWA.

67 + ``flash_attn_version`` is forced to 4 if not already set (R-SWA

68 + mask_mod requires FA4; FA3 cannot express it). Raises if FA4 is

69 + not available on this device.

70 +

71 + 2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'``

72 + → FlexAttention R-SWA via Triton block mask.

73 +

74 + 3. ``--attention-config '{"backend": "auto"}'`` (or omitted)

75 + → Auto-detect: FA4 if available (H20/H100 SM90), else FlexAttention.

76 +

77 + Regardless of backend, prefix caching is disabled for this model: R-SWA

78 + decode-phase KV is not a pure causal function of the prefix (so decode

79 + blocks are not reusable), and single-turn image-led OCR prompts rarely

80 + hit the prefix cache.

81 +

82 + Example — force FlexAttention even on a machine with FA4::

83 +

84 + vllm serve baidu/Unlimited-OCR \\

85 + --attention-config '{"backend": "FLEX_ATTENTION"}'

86 + """

87 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

88 + from vllm.vllm_flash_attn import is_fa_version_supported

89 +

90 + attn_config = vllm_config.attention_config

91 + fa4_available = is_fa_version_supported(4)

92 +

93 + # ── step 1: resolve backend ─────────────────────────────────────────

94 + # None means the user did not explicitly specify a backend; auto-select.

95 + if attn_config.backend is None:

96 + attn_config.backend = (

97 + AttentionBackendEnum.FLASH_ATTN

98 + if fa4_available

99 + else AttentionBackendEnum.FLEX_ATTENTION

100 + )

101 + logger.info(

102 + "Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).",

103 + attn_config.backend.value,

104 + fa4_available,

105 + )

106 +

107 + # ── step 2: configure the chosen backend ────────────────────────────

108 + if attn_config.backend == AttentionBackendEnum.FLASH_ATTN:

109 + if not fa4_available:

110 + raise RuntimeError(

111 + "Unlimited-OCR: --attention-config backend=FLASH_ATTN "

112 + "requires FA4 (rswa_mask_mod), but FA4 is not available on "

113 + "this device/installation. Use backend=FLEX_ATTENTION or "

114 + "upgrade vllm-flash-attn."

115 + )

116 + # On SM90 (H20), the default FA version is FA3 regardless of FA4

117 + # availability (FA4 is only auto-upgraded when head_size > 256).

118 + # The R-SWA mask_mod requires FA4, so force the version globally.

119 + if attn_config.flash_attn_version is None:

120 + attn_config.flash_attn_version = 4

121 + elif attn_config.flash_attn_version < 4:

122 + logger.warning(

123 + "Unlimited-OCR: flash_attn_version=%d cannot express the "

124 + "R-SWA mask_mod; upgrading to 4.",

125 + attn_config.flash_attn_version,

126 + )

127 + attn_config.flash_attn_version = 4

128 + logger.info(

129 + "Unlimited-OCR: FlashAttention FA%d + rswa_mask_mod — exact R-SWA.",

130 + attn_config.flash_attn_version,

131 + )

132 +

133 + elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION:

134 + logger.info(

135 + "Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.",

136 + ""

137 + if not fa4_available

138 + else (

139 + " (FA4 available but not used; pass backend=FLASH_ATTN to upgrade)"

140 + ),

141 + )

142 +

143 + else:

144 + raise ValueError(

145 + f"Unlimited-OCR: unsupported attention backend "

146 + f"{attn_config.backend!r} for R-SWA. "

147 + "Use FLASH_ATTN (FA4) or FLEX_ATTENTION."

148 + )

149 +

150 + # R-SWA windows the *generated* tokens, so a decode-token's KV is not a

151 + # pure causal function of the prefix and cannot be safely reused across

152 + # requests via prefix caching. Only the prompt/image prefix is cacheable,

153 + # but OCR is single-turn with image-led prompts that rarely share a

154 + # prefix, so prefix caching brings little benefit while complicating the

155 + # KV cache manager. Disable it for this model.

156 + cache_config = vllm_config.cache_config

157 + if cache_config.enable_prefix_caching:

158 + cache_config.enable_prefix_caching = False

159 + logger.info(

160 + "Unlimited-OCR: disabling prefix caching (R-SWA decode KV is not "

161 + "cacheable, and single-turn image-led prompts rarely hit the "

162 + "prefix cache)."

163 + )

164 +

165 + mm_config = getattr(vllm_config.model_config, "multimodal_config", None)

166 + if mm_config is not None:

167 + if mm_config.mm_encoder_attn_backend is None:

168 + mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN

169 + elif mm_config.mm_encoder_attn_backend == AttentionBackendEnum.FLASHINFER:

170 + logger.warning(

171 + "Unlimited-OCR: FlashInfer is not supported for the vision "

172 + "encoder (the CLIP stage runs full attention without "

173 + "cu_seqlens); falling back to FlashAttention."

174 + )

175 + mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN

176 +

177 + @staticmethod

178 + def verify_and_update_model_config(model_config: "ModelConfig") -> None:

179 + text_config = model_config.hf_config.text_config

180 + text_config.architectures = ["DeepseekV2ForCausalLM"]

181 + if getattr(model_config.hf_config, "rswa_window", None) is None:

182 + model_config.hf_config.rswa_window = 128

183 + # Propagate rswa_window to text_config so that DeepseekAttention (which

184 + # receives text_config as its vllm_config.model_config.hf_config via

185 + # init_vllm_registered_model) can read it and create RSWAAttention.

186 + rswa_window = model_config.hf_config.rswa_window

187 + text_config.rswa_window = rswa_window

188 +

189 +

190 class Gemma4Config(VerifyAndUpdateConfig):

191 @staticmethod

192 def verify_and_update_config(vllm_config: "VllmConfig") -> None:

@@ -703,6 +836,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {

836 "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,

837 "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,

838 "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,

839 + "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig,

840 "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,

841 "XLMRobertaModel": JinaRobertaModelConfig,

842 }

vllm/model_executor/models/deepseek_ocr.py

+4 -2

@@ -220,8 +220,10 @@ class DeepseekOCRProcessingInfo(BaseProcessingInfo):

220 patch_size = 16

221 downsample_ratio = 4

222

223 - if CROP_MODE:

224 - if image_width <= 640 and image_height <= 640:

223 + # Use the caller-supplied `cropping` flag so that callers that disable

224 + # crop mode for multi-image requests get a consistent token count.

225 + if cropping:

226 + if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE:

227 crop_ratio = [1, 1]

228 else:

229 # find the closest aspect ratio to the target

vllm/model_executor/models/deepseek_v2.py

+25 -12

@@ -45,7 +45,7 @@ from vllm.distributed import (

45 )

46 from vllm.logger import init_logger

47 from vllm.model_executor.layers.activation import SiluAndMul

48 -from vllm.model_executor.layers.attention import Attention

48 +from vllm.model_executor.layers.attention import Attention, RSWAAttention

49 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

50 from vllm.model_executor.layers.fused_moe import (

51 FusedMoE,

@@ -174,15 +174,28 @@ class DeepseekAttention(nn.Module):

174 max_position=max_position_embeddings,

175 rope_parameters=config.rope_parameters,

176 )

177 - self.attn = Attention(

178 - self.num_heads,

179 - self.head_dim,

180 - self.scaling,

181 - num_kv_heads=self.num_kv_heads,

182 - cache_config=cache_config,

183 - quant_config=quant_config,

184 - prefix=f"{prefix}.attn",

185 - )

177 + rswa_window = getattr(vllm_config.model_config.hf_config, "rswa_window", None)

178 + if rswa_window is not None:

179 + self.attn = RSWAAttention(

180 + self.num_heads,

181 + self.head_dim,

182 + self.scaling,

183 + num_kv_heads=self.num_kv_heads,

184 + cache_config=cache_config,

185 + quant_config=quant_config,

186 + prefix=f"{prefix}.attn",

187 + rswa_window=rswa_window,

188 + )

189 + else:

190 + self.attn = Attention(

191 + self.num_heads,

192 + self.head_dim,

193 + self.scaling,

194 + num_kv_heads=self.num_kv_heads,

195 + cache_config=cache_config,

196 + quant_config=quant_config,

197 + prefix=f"{prefix}.attn",

198 + )

199

200 def forward(

201 self,

@@ -588,12 +601,12 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):

601 compilation_config.static_forward_context[prefix] = self

602

603 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

591 - return MLAAttentionSpec( # Only has one vector instead of K + V

604 + return MLAAttentionSpec(

605 block_size=self.cache_config.block_size,

606 num_kv_heads=1,

607 head_size=self.head_dim,

608 dtype=self.dtype,

596 - )

609 + ) # Only has one vector instead of K + V

610

611 def forward(self): ...

612

vllm/model_executor/models/registry.py

@@ -358,6 +358,7 @@ _MULTIMODAL_MODELS = {

358 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),

359 "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),

360 "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),

361 + "UnlimitedOCRForCausalLM": ("unlimited_ocr", "UnlimitedOCRForCausalLM"),

362 "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),

363 "Eagle2_5_VLForConditionalGeneration": (

364 "eagle2_5_vl",

vllm/model_executor/models/unlimited_ocr.py new

+250

@@ -0,0 +1,250 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Inference-only Unlimited-OCR model compatible with HuggingFace weights.

4 +

5 +Unlimited-OCR (``baidu/Unlimited-OCR``) shares

6 +the exact DeepSeek-OCR (gundam, ``base_size=1024`` / ``image_size=640`` / crop)

7 +vision stack: a DeepEncoder (SAM-ViT-B + CLIP-L) followed by a linear MLP

8 +projector, with the same image-token tiling layout. The only difference is the

9 +language backbone, which is a DeepSeek-V2 *MoE* (64 routed + 2 shared experts,

10 +``first_k_dense_replace=1``) that uses plain multi-head attention

11 +(``use_mla=False``, ``qk_nope_head_dim == qk_rope_head_dim == 0``) instead of

12 +the dense MLA decoder used by DeepSeek-OCR.

13 +

14 +vLLM's ``DeepseekV2DecoderLayer`` already dispatches to the plain-MHA

15 +``DeepseekAttention`` whenever ``qk_nope_head_dim == qk_rope_head_dim == 0`` and

16 +builds the MoE blocks straight from the config, so the whole DeepSeek-OCR

17 +multimodal wrapper can be reused verbatim. Model-specific config (language

18 +backbone architecture, FlexAttention for R-SWA, vision encoder backend, and

19 +``rswa_window``) is applied in ``UnlimitedOCRForCausalLMConfig``.

20 +

21 +Attention backend: the reference applies Reference Sliding Window Attention

22 +(R-SWA) -- the prompt/image tokens form a globally-visible prefix while the

23 +*generated* tokens additionally attend only a fixed sliding window (128) of

24 +recent tokens. We reproduce this (Level 1: full KV cache + custom mask) by

25 +forcing the language model onto the FlexAttention backend and installing an

26 +R-SWA ``mask_mod``. FlexAttention is the only backend able to express the

27 +"global prefix + sliding window" mask; FlashAttention-3 / Triton only support a

28 +uniform window (and additionally crash or compute incorrectly on this decoder's

29 +10-head,

30 +head_dim-128 shape), and FlashInfer's paged decode exposes no custom mask. The

31 +window size is published via ``model_config.rswa_window``, which the model

32 +runner reads to plumb per-request prefix lengths into the FlexAttention mask.

33 +

34 +The *vision encoder* (DeepEncoder's CLIP stage, head_dim 64) is unaffected and

35 +does not use R-SWA: it runs a single full-attention prefill pass. FlashAttention,

36 +Triton and torch SDPA all produce correct, equally fast results; only FlashInfer

37 +is incompatible (its ViT path asserts on the varlen cu_seqlens metadata that this

38 +CLIP encoder never builds). We default the encoder to FlashAttention and

39 +transparently fall back to it if FlashInfer is requested.

40 +

41 +To suppress repetition on long documents, use ``NGramPerReqLogitsProcessor`` from

42 +this module (same request-level processor as DeepSeek-OCR) with::

43 +

44 + SamplingParams(

45 + temperature=0.0,

46 + max_tokens=8192,

47 + extra_args={"ngram_size": 35, "window_size": 128},

48 + )

49 +

50 +Image processing

51 +----------------

52 +Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR), i.e.

53 +``dynamic_preprocess`` runs with ``max_num=32``.

54 +

55 +Multi-image requests fall back to non-crop mode: crop ("gundam") mode is only

56 +used for single-image input. DeepSeek-OCR does *not* have this restriction.

57 +

58 +Because that fallback makes the per-image processor output depend on *how many*

59 +images are in the request, it breaks the assumption behind vLLM's per-item

60 +multimodal processing cache (``MultiModalProcessorOnlyCache``). We handle this

61 +the same way ``DeepseekVL2MultiModalProcessor`` does: only the single-image case

62 +(which always crops) is cached, while multi-image requests bypass the cache and

63 +are recomputed fresh -- see ``_cached_apply_hf_processor`` below. This keeps the

64 +processing cache consistent (verified by ``test_processing_correctness``).

65 +"""

66 +

67 +import math

68 +from collections.abc import Mapping, Sequence

69 +

70 +from vllm.config import VllmConfig

71 +from vllm.multimodal import MULTIMODAL_REGISTRY

72 +from vllm.multimodal.inputs import MultiModalKwargsItems

73 +from vllm.multimodal.parse import (

74 + ImageEmbeddingItems,

75 + ImageProcessorItems,

76 + ImageSize,

77 + MultiModalDataItems,

78 +)

79 +from vllm.multimodal.processing import PromptReplacement, PromptUpdate

80 +from vllm.multimodal.processing.context import TimingContext

81 +from vllm.multimodal.processing.inputs import ProcessorInputs

82 +from vllm.multimodal.processing.processor import MultiModalProcessingInfo

83 +from vllm.transformers_utils.processors.deepseek_ocr import (

84 + BASE_SIZE,

85 + CROP_MODE,

86 + IMAGE_SIZE,

87 + count_tiles,

88 +)

89 +

90 +from .deepseek_ocr import (

91 + DeepseekOCRDummyInputsBuilder,

92 + DeepseekOCRForCausalLM,

93 + DeepseekOCRMultiModalProcessor,

94 + DeepseekOCRProcessingInfo,

95 + NGramPerReqLogitsProcessor,

96 +)

97 +

98 +__all__ = [

99 + "NGramPerReqLogitsProcessor",

100 + "UnlimitedOCRForCausalLM",

101 +]

102 +

103 +# Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR).

104 +_UNLIMITED_OCR_MAX_CROPS = 32

105 +

106 +

107 +class UnlimitedOCRProcessingInfo(DeepseekOCRProcessingInfo):

108 + """ProcessingInfo for Unlimited-OCR: same as DeepSeek-OCR but with

109 + max_crops=32 instead of 6. The higher crop count allows tiling very large

110 + document pages into up to 32 640×640 patches (dynamic_preprocess max_num=32).

111 + """

112 +

113 + def get_hf_config(self):

114 + from vllm.transformers_utils.configs.unlimited_ocr import UnlimitedOCRConfig

115 +

116 + return self.ctx.get_hf_config(UnlimitedOCRConfig)

117 +

118 + def get_hf_processor(self, **kwargs: object):

119 + from vllm.transformers_utils.processors.unlimited_ocr import (

120 + UnlimitedOCRProcessor,

121 + )

122 +

123 + v1_processor_config = dict(

124 + image_size=IMAGE_SIZE,

125 + base_size=BASE_SIZE,

126 + crop_mode=CROP_MODE,

127 + strategy="v1",

128 + max_crops=_UNLIMITED_OCR_MAX_CROPS,

129 + )

130 + return self.ctx.get_hf_processor(

131 + UnlimitedOCRProcessor,

132 + **{**v1_processor_config, **kwargs},

133 + )

134 +

135 + def get_num_image_tokens(

136 + self, *, image_width: int, image_height: int, cropping: bool = True

137 + ) -> int:

138 + patch_size = 16

139 + downsample_ratio = 4

140 +

141 + # Honour the caller-supplied `cropping` flag: multi-image callers pass

142 + # cropping=False to match UnlimitedOCRProcessor.tokenize_with_images.

143 + if cropping:

144 + if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE:

145 + crop_ratio = [1, 1]

146 + else:

147 + crop_ratio = count_tiles(

148 + image_width,

149 + image_height,

150 + max_num=_UNLIMITED_OCR_MAX_CROPS,

151 + image_size=IMAGE_SIZE,

152 + )

153 + num_width_tiles, num_height_tiles = crop_ratio

154 + else:

155 + num_width_tiles = num_height_tiles = 1

156 +

157 + h = w = math.ceil((BASE_SIZE // patch_size) / downsample_ratio)

158 + h2 = w2 = math.ceil((IMAGE_SIZE // patch_size) / downsample_ratio)

159 +

160 + global_views_tokens = h * (w + 1)

161 + if num_width_tiles > 1 or num_height_tiles > 1:

162 + local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1)

163 + else:

164 + local_views_tokens = 0

165 +

166 + return global_views_tokens + local_views_tokens + 1

167 +

168 + def get_image_size_with_most_features(self) -> ImageSize:

169 + # With max_crops=32, the widest possible grid is 4×8 (aspect ratio 1:2).

170 + # A 2560×5120 image (4×640 × 8×640) selects exactly 4×8=32 tiles and

171 + # produces the maximum token count.

172 + return ImageSize(width=640 * 4, height=640 * 8)

173 +

174 +

175 +class UnlimitedOCRMultiModalProcessor(DeepseekOCRMultiModalProcessor):

176 + """Multimodal processor for Unlimited-OCR.

177 +

178 + Disables crop mode for multi-image requests (to stay consistent with

179 + ``UnlimitedOCRProcessor.tokenize_with_images``), and -- since that makes the

180 + per-image output depend on the request's image count -- bypasses the

181 + per-item processing cache for multi-image requests, exactly like

182 + ``DeepseekVL2MultiModalProcessor``.

183 +

184 + DeepSeek-OCR does *not* apply either of these.

185 + """

186 +

187 + def _get_prompt_updates(

188 + self,

189 + mm_items: MultiModalDataItems,

190 + hf_processor_mm_kwargs: Mapping[str, object],

191 + out_mm_kwargs: MultiModalKwargsItems,

192 + ) -> Sequence[PromptUpdate]:

193 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

194 +

195 + image_token_id = hf_processor.image_token_id

196 + assert isinstance(image_token_id, int)

197 +

198 + def get_replacement_unlimited_ocr(item_idx: int):

199 + images = mm_items.get_items(

200 + "image", (ImageEmbeddingItems, ImageProcessorItems)

201 + )

202 +

203 + if isinstance(images, ImageEmbeddingItems):

204 + num_image_tokens = images.get_feature_size(item_idx)

205 + else:

206 + size = images.get_image_size(item_idx)

207 +

208 + # Disable crop mode for multi-image input.

209 + # UnlimitedOCRProcessor.tokenize_with_images applies the same

210 + # fallback, so both paths must agree on the effective crop flag.

211 + effective_cropping = CROP_MODE and len(images) == 1

212 +

213 + num_image_tokens = self.info.get_num_image_tokens(

214 + image_width=size.width,

215 + image_height=size.height,

216 + cropping=effective_cropping,

217 + )

218 + return [image_token_id] * num_image_tokens

219 +

220 + return [

221 + PromptReplacement(

222 + modality="image",

223 + target=[image_token_id],

224 + replacement=get_replacement_unlimited_ocr,

225 + )

226 + ]

227 +

228 + def _cached_apply_hf_processor(

229 + self,

230 + inputs: ProcessorInputs,

231 + timing_ctx: TimingContext,

232 + ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

233 + # The processor logic differs for single-image (crop) vs multi-image

234 + # (no crop) requests. The processing cache assumes per-item output is

235 + # invariant of how many images are passed per prompt, so we only cache

236 + # the single-image case and recompute multi-image requests fresh.

237 + if inputs.mm_data_items.get_count("image", strict=False) > 1:

238 + return self._apply_hf_processor(inputs, timing_ctx)

239 +

240 + return super()._cached_apply_hf_processor(inputs, timing_ctx)

241 +

242 +

243 +@MULTIMODAL_REGISTRY.register_processor(

244 + UnlimitedOCRMultiModalProcessor,

245 + info=UnlimitedOCRProcessingInfo,

246 + dummy_inputs=DeepseekOCRDummyInputsBuilder,

247 +)

248 +class UnlimitedOCRForCausalLM(DeepseekOCRForCausalLM):

249 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

250 + super().__init__(vllm_config=vllm_config, prefix=prefix)

vllm/tokenizers/registry.py

+5 -1

@@ -31,7 +31,11 @@ logger = init_logger(__name__)

31 # temporary workaround and better long term solutions are:

32 # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better)

33 # - Fix tokenizer_class on the hub for the affected models (best)

34 -_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"}

34 +_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {

35 + "step3_vl",

36 + "step3p7",

37 + "unlimited-ocr",

38 +}

40 _VLLM_TOKENIZERS = {

41 "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] = {

29 "colpali": CHAT_TEMPLATES_DIR / "template_basic.jinja",

30 "deepseek_ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",

31 "deepseek_ocr2": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",

32 + "unlimited-ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",

33 "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja",

34 "fuyu": CHAT_TEMPLATES_DIR / "template_fuyu.jinja",

35 "minicpmv": _get_minicpmv_chat_template_fallback,

vllm/transformers_utils/config.py

@@ -124,6 +124,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

124 laguna="LagunaConfig",

125 lfm2_moe="Lfm2MoeConfig",

126 tarsier2="Tarsier2Config",

127 + **{"unlimited-ocr": "UnlimitedOCRConfig"},

128 )

129

130 _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"}

vllm/transformers_utils/configs/__init__.py

@@ -73,6 +73,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

73 "RadioConfig": "vllm.transformers_utils.configs.radio",

74 "SpeculatorsConfig": "vllm.transformers_utils.configs.speculators",

75 "UltravoxConfig": "vllm.transformers_utils.configs.ultravox",

76 + "UnlimitedOCRConfig": "vllm.transformers_utils.configs.unlimited_ocr",

77 "Step3VLConfig": "vllm.transformers_utils.configs.step3_vl",

78 "Step3VisionEncoderConfig": "vllm.transformers_utils.configs.step3_vl",

79 "Step3TextConfig": "vllm.transformers_utils.configs.step3_vl",

@@ -147,6 +148,7 @@ __all__ = [

148 "RadioConfig",

149 "SpeculatorsConfig",

150 "UltravoxConfig",

151 + "UnlimitedOCRConfig",

152 "Step3VLConfig",

153 "Step3VisionEncoderConfig",

154 "Step3TextConfig",

vllm/transformers_utils/configs/unlimited_ocr.py new

+35

@@ -0,0 +1,35 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +# Unlimited-OCR (baidu/Unlimited-OCR) reuses

5 +# the DeepSeek-OCR multimodal layout (DeepEncoder = SAM-ViT-B + CLIP-L, a linear

6 +# MLP projector and a DeepSeek-V2 text backbone). The only architectural

7 +# difference is the language model, which is a DeepSeek-V2 *MoE* with plain

8 +# multi-head attention (``use_mla=False``) instead of the dense MLA backbone.

9 +# We therefore reuse ``DeepseekVLV2Config`` for parsing the nested config.

10 +

11 +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config

12 +

13 +

14 +class UnlimitedOCRConfig(DeepseekVLV2Config):

15 + model_type = "unlimited-ocr"

16 +

17 + # An explicit ``__init__`` is required: Transformers v5 processes each

18 + # concrete config class' ``__init__`` signature to build nested sub-configs,

19 + # and an empty subclass (only overriding ``model_type``) would skip

20 + # ``DeepseekVLV2Config.__init__``, leaving ``text_config`` unset.

21 + def __init__(

22 + self,

23 + tile_tag: str = "2D",

24 + global_view_pos: str = "head",

25 + candidate_resolutions: tuple[tuple[int, int]] = ((384, 384),),

26 + rswa_window: int = 128,

27 + **kwargs,

28 + ):

29 + super().__init__(

30 + tile_tag=tile_tag,

31 + global_view_pos=global_view_pos,

32 + candidate_resolutions=candidate_resolutions,

33 + **kwargs,

34 + )

35 + self.rswa_window = rswa_window

vllm/transformers_utils/model_arch_config_convertor.py

@@ -310,6 +310,12 @@ class ModelArchConfigConvertorBase:

310 return False

311 return self.hf_config.model_type in MM_PREFIX_LM_MODELS

312

313 + def rswa_window(self) -> int | None:

314 + value = getattr(self.hf_config, "rswa_window", None)

315 + if value is None:

316 + return None

317 + return int(value)

318 +

319 def derive_max_model_len_and_key(self) -> tuple[float, str | None]:

320 derived_max_model_len = float("inf")

321 possible_keys = [

@@ -360,6 +366,7 @@ class ModelArchConfigConvertorBase:

366 quantization_config=self.get_quantization_config(),

367 is_deepseek_mla=self.is_deepseek_mla(),

368 is_mm_prefix_lm=self.is_mm_prefix_lm(),

369 + rswa_window=self.rswa_window(),

370 derived_max_model_len_and_key=self.derive_max_model_len_and_key(),

371 )

372

vllm/transformers_utils/processors/deepseek_ocr.py

+3 -1

@@ -161,10 +161,12 @@ class DeepseekOCRProcessor(ProcessorMixin):

161 image_size: int = IMAGE_SIZE,

162 base_size: int = BASE_SIZE,

163 strategy: Literal["v1", "v2"] = "v1",

164 + max_crops: int = MAX_CROPS,

165 **kwargs,

166 ):

167 self.image_size = image_size

168 self.base_size = base_size

169 + self.max_crops = max_crops

170

171 # image token calculation strategy for

172 # Deepseek-OCR and Deepseek-OCR-2

@@ -332,7 +334,7 @@ class DeepseekOCRProcessor(ProcessorMixin):

334 crop_ratio = [1, 1]

335 elif cropping:

336 images_crop_raw, crop_ratio = dynamic_preprocess(

335 - image, image_size=self.image_size

337 + image, image_size=self.image_size, max_num=self.max_crops

338 )

339 else:

340 crop_ratio = [1, 1]

vllm/transformers_utils/processors/unlimited_ocr.py new

+46

@@ -0,0 +1,46 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Image processor for Unlimited-OCR (baidu/Unlimited-OCR)."""

4 +

5 +from PIL import Image

6 +

7 +from vllm.logger import init_logger

8 +from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor

9 +

10 +logger = init_logger(__name__)

11 +

12 +

13 +class UnlimitedOCRProcessor(DeepseekOCRProcessor):

14 + """DeepseekOCRProcessor variant for Unlimited-OCR.

15 +

16 + The only behavioural difference from the base processor is a multi-image

17 + safeguard: when more than one image is present, crop ("gundam") mode is

18 + disabled.

19 +

20 + Because the effective crop flag then depends on *how many* images are in the

21 + request, the per-item processing output is no longer invariant of sibling

22 + images. ``UnlimitedOCRMultiModalProcessor`` accounts for this by bypassing

23 + the multimodal processing cache for multi-image requests (see its

24 + ``_cached_apply_hf_processor``), so the two paths stay consistent.

25 +

26 + DeepSeek-OCR does *not* have this restriction because its ``max_crops=6`` is

27 + small enough to be safe for multi-image use.

28 + """

29 +

30 + def tokenize_with_images(

31 + self,

32 + conversation: str,

33 + images: list[Image.Image],

34 + bos: bool = True,

35 + eos: bool = True,

36 + cropping: bool = True,

37 + ):

38 + if len(images) > 1 and cropping:

39 + logger.warning_once(

40 + "Unlimited-OCR: crop mode is not supported for multi-image "

41 + "input. Falling back to cropping=False."

42 + )

43 + cropping = False

44 + return super().tokenize_with_images(

45 + conversation, images, bos=bos, eos=eos, cropping=cropping

46 + )

vllm/v1/attention/backend.py

@@ -455,6 +455,13 @@ class CommonAttentionMetadata:

455 where bidirectional attention should apply. None for text-only

456 batches or non-PrefixLM models."""

457

458 + rswa_prefix_lens: torch.Tensor | None = None

459 + """(batch_size,) per-request prefix length (prompt/image token count) for

460 + Reference Sliding Window Attention (R-SWA). Tokens with logical index below

461 + this stay globally visible; later (generated) tokens additionally see a

462 + fixed sliding window. None disables R-SWA. The attention backend copies this

463 + into its own persistent buffer and reads ``rswa_window`` from model config."""

464 +

465 # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)

466 _seq_lens_cpu: torch.Tensor | None = None

467 _num_computed_tokens_cpu: torch.Tensor | None = None

@@ -539,6 +546,7 @@ class CommonAttentionMetadata:

546 dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),

547 dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),

548 is_prefilling=maybe_slice_reqs(self.is_prefilling),

549 + rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),

550 )

551

552

vllm/v1/attention/backends/flash_attn.py

+114 -3

@@ -256,6 +256,16 @@ class FlashAttentionMetadata:

256 # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.

257 mm_prefix_range_tensor: torch.Tensor | None = None

258

259 + # Reference Sliding Window Attention (R-SWA) fields.

260 + # rswa_prefix_lens: per-request prompt lengths [num_reqs], int32, CUDA.

261 + # rswa_window: sliding window size (scalar int, for logic checks).

262 + # rswa_window_tensor: [1] int32 CUDA tensor — pre-allocated in build() so

263 + # no CPU→CUDA copy is needed inside forward() during CUDA graph capture.

264 + # Only populated when the model uses R-SWA (Unlimited-OCR).

265 + rswa_prefix_lens: torch.Tensor | None = None

266 + rswa_window: int | None = None

267 + rswa_window_tensor: torch.Tensor | None = None

268 +

269

270 def _get_sliding_window_configs(

271 vllm_config: VllmConfig,

@@ -386,6 +396,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

396 # populated on first build() call.

397 self.aot_sliding_window: tuple[int, int] | None = None

398

399 + # R-SWA: persistent CUDA-graph-safe buffers owned by this builder.

400 + self.rswa_window: int | None = self.model_config.rswa_window

401 + self.persistent_rswa_prefix_lens: torch.Tensor | None = None

402 + self.persistent_rswa_window_tensor: torch.Tensor | None = None

403 + if self.rswa_window is not None:

404 + max_num_reqs = vllm_config.scheduler_config.max_num_seqs

405 + self.persistent_rswa_prefix_lens = torch.zeros(

406 + max_num_reqs, dtype=torch.int32, device=self.device

407 + )

408 + self.persistent_rswa_window_tensor = torch.tensor(

409 + [self.rswa_window], dtype=torch.int32, device=self.device

410 + )

411 +

412 def build(

413 self,

414 common_prefix_len: int,

@@ -589,6 +612,22 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

612 mm_ranges, num_reqs, seq_lens.device

613 )

614

615 + # R-SWA: copy prefix lengths into persistent buffers (outside the

616 + # compiled region) so forward() never allocates during CUDA graph

617 + # capture. rswa_window is a static model config scalar read here.

618 + if (

619 + self.rswa_window is not None

620 + and common_attn_metadata.rswa_prefix_lens is not None

621 + ):

622 + assert self.persistent_rswa_prefix_lens is not None

623 + assert self.persistent_rswa_window_tensor is not None

624 + src = common_attn_metadata.rswa_prefix_lens

625 + rswa_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]

626 + rswa_prefix_lens.copy_(src[:num_reqs], non_blocking=True)

627 + attn_metadata.rswa_prefix_lens = rswa_prefix_lens

628 + attn_metadata.rswa_window = self.rswa_window

629 + attn_metadata.rswa_window_tensor = self.persistent_rswa_window_tensor

630 +

631 return attn_metadata

632

633 def update_block_table(

@@ -805,7 +844,7 @@ class FlashAttentionImpl(AttentionImpl):

844 )

845 return output

846 else:

808 - sliding_window_size = (

847 + sliding_window_size: list[int] | None = (

848 list(self.sliding_window)

849 if self.sliding_window is not None

850 else None

@@ -840,6 +879,25 @@ class FlashAttentionImpl(AttentionImpl):

879 mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges)

880 mm_aux = [mm_prefix_ranges]

881

882 + # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level

883 + # mask without block-size approximation. The mask_mod encodes

884 + # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which

885 + # supersedes any FA-layer sliding_window_size parameter.

886 + rswa_mask_mod_fn = None

887 + rswa_aux = None

888 + if (

889 + attn_metadata.rswa_prefix_lens is not None

890 + and self.vllm_flash_attn_version == 4

891 + and not is_dynamic_causal

892 + ):

893 + rswa_mask_mod_fn = _make_rswa_mask_mod()

894 + rswa_aux = [

895 + attn_metadata.rswa_prefix_lens.to(torch.int32),

896 + attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor

897 + ]

898 + # mask_mod fully expresses R-SWA; disable FA's own window.

899 + sliding_window_size = None

900 +

901 dynamic_causal = None

902 if isinstance(causal, torch.Tensor):

903 if self.vllm_flash_attn_version != 4:

@@ -873,8 +931,8 @@ class FlashAttentionImpl(AttentionImpl):

931 dynamic_causal=dynamic_causal,

932 num_splits=attn_metadata.max_num_splits,

933 s_aux=self.sinks,

876 - mask_mod=mm_mask_mod,

877 - aux_tensors=mm_aux,

934 + mask_mod=rswa_mask_mod_fn or mm_mask_mod,

935 + aux_tensors=rswa_aux or mm_aux,

936 )

937 return output

938

@@ -1152,6 +1210,59 @@ def _make_mm_prefix_mask_mod(max_ranges: int):

1210 return mm_prefix_mask_mod

1211

1212

1213 +def _make_rswa_mask_mod():

1214 + """Build a CuTE-DSL mask_mod for Reference Sliding Window Attention (R-SWA).

1215 +

1216 + FA4 varlen + paged-KV convention (verified from cute/mask.py apply_mask):

1217 + q_idx = LOCAL query-token offset (0 .. seqlen_q - 1) within this sequence.

1218 + kv_idx = LOCAL KV-token position (0 .. seqlen_k - 1) within this sequence.

1219 +

1220 + To recover the ABSOLUTE token position (needed for causal and the sliding

1221 + window distance), use the standard offset:

1222 + abs_q = q_idx + (seqlen_k - seqlen_q)

1223 +

1224 + R-SWA keep condition:

1225 + abs_q >= kv_idx (causal: KV at or before the query)

1226 + AND (kv_idx < prefix_len (global prefix is always visible)

1227 + OR abs_q - kv_idx < window) (generated tokens: sliding window)

1228 +

1229 + aux_tensors[0]: prefix_lens [num_reqs] int32 — per-request prefill length.

1230 + aux_tensors[1]: rswa_window [1] int32 — decode sliding window size.

1231 +

1232 + use_fast_sampling=True lets FA4 skip fully-masked KV blocks (gap blocks)

1233 + without loading their data.

1234 + """

1235 + import cutlass.cute as cute

1236 + from cutlass import Int32 # type: ignore[attr-defined]

1237 +

1238 + from vllm.vllm_flash_attn.cute.utils import ( # type: ignore[import-untyped]

1239 + scalar_to_ssa,

1240 + )

1241 +

1242 + @cute.jit

1243 + def rswa_mask_mod(

1244 + batch_idx: cute.TensorSSA,

1245 + head_idx: cute.TensorSSA,

1246 + q_idx: cute.TensorSSA,

1247 + kv_idx: cute.TensorSSA,

1248 + seqlen_info,

1249 + aux_tensors,

1250 + ):

1251 + b = batch_idx[0]

1252 + prefix_len = scalar_to_ssa(aux_tensors[0][b], Int32)

1253 + window = scalar_to_ssa(aux_tensors[1][0], Int32)

1254 + # Convert local q offset to absolute token position.

1255 + offset = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)

1256 + abs_q = q_idx + offset

1257 + causal = kv_idx <= abs_q

1258 + in_prefix = kv_idx < prefix_len

1259 + in_window = (abs_q - kv_idx) < window

1260 + return causal & (in_prefix | in_window)

1261 +

1262 + rswa_mask_mod.use_fast_sampling = True

1263 + return rswa_mask_mod

1264 +

1265 +

1266 def use_cascade_attention(

1267 common_prefix_len: int,

1268 query_lens: np.ndarray,

vllm/v1/attention/backends/flex_attention.py

+148 -4

@@ -408,6 +408,11 @@ class FlexAttentionMetadata:

408 sliding_window: int | None = None

409 mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None

410 block_sparsity_hint: BlockSparsityHint | None = None

411 + # Reference Sliding Window Attention (R-SWA): per-request prefix length

412 + # (prompt/image tokens stay globally visible) plus a sliding window over

413 + # generated tokens. Both must be set to enable.

414 + rswa_prefix_lens: torch.Tensor | None = None

415 + rswa_window: int | None = None

416

417 @cached_property

418 def logical_block_ids(self):

@@ -571,6 +576,52 @@ class FlexAttentionMetadata:

576

577 return final_mask_mod

578

579 + def get_rswa_mask_mod(self) -> _mask_mod_signature:

580 + """Creates the Reference Sliding Window Attention (R-SWA) mask_mod.

581 +

582 + R-SWA keeps the whole prefix (image + prompt tokens, i.e. logical index

583 + ``< prefix_len``) globally visible while generated tokens additionally

584 + attend a fixed sliding window of recent tokens. This term is combined

585 + with the base causal mask via logical AND, so it only ever *removes*

586 + far-away generated tokens that fall outside the window and outside the

587 + prefix.

588 + """

589 +

590 + assert self.doc_ids is not None

591 + assert self.rswa_prefix_lens is not None

592 + assert self.rswa_window is not None

593 + doc_ids = self.doc_ids

594 + prefix_lens = self.rswa_prefix_lens

595 + window = self.rswa_window

596 +

597 + def rswa_mask_mod(

598 + q_req: torch.Tensor,

599 + logical_q_idx: torch.Tensor,

600 + logical_kv_idx: torch.Tensor,

601 + ) -> torch.Tensor:

602 + prefix_len = prefix_lens[q_req]

603 + in_prefix = logical_kv_idx < prefix_len

604 + in_window = (logical_q_idx - logical_kv_idx) < window

605 + return in_prefix | in_window

606 +

607 + def final_mask_mod(

608 + b: torch.Tensor,

609 + h: torch.Tensor,

610 + q_idx: torch.Tensor,

611 + physical_kv_idx: torch.Tensor,

612 + ) -> torch.Tensor:

613 + (is_valid, logical_q_idx, logical_kv_idx) = (

614 + self._convert_physical_to_logical(doc_ids, q_idx, physical_kv_idx)

615 + )

616 + q_req = doc_ids[q_idx]

617 + return torch.where(

618 + is_valid,

619 + rswa_mask_mod(q_req, logical_q_idx, logical_kv_idx),

620 + False,

621 + )

622 +

623 + return final_mask_mod

624 +

625 def get_mask_mod(self):

626 # Stage-1: initialize the base mask_mod

627 # (causal mask for decoder or bidirectional mask for encoder)

@@ -588,6 +639,10 @@ class FlexAttentionMetadata:

639 # Add prefix LM mask for vision-language prefix LM attention

640 prefix_lm_mask_mod = self.get_prefix_lm_mask_mod()

641 mask_mod = or_masks(mask_mod, prefix_lm_mask_mod)

642 + if self.rswa_window is not None and self.rswa_prefix_lens is not None:

643 + # Reference Sliding Window Attention: AND with the base causal mask

644 + # (prefix stays global, generated tokens use a sliding window).

645 + mask_mod = and_masks(mask_mod, self.get_rswa_mask_mod())

646 return mask_mod

647

648 def get_transformed_score_mod(self) -> _score_mod_signature | None:

@@ -663,9 +718,21 @@ class FlexAttentionMetadata:

718 self.doc_ids, : cdiv(self.max_seq_len, self.block_size)

719 ]

720

721 + # block_table slots beyond each request's seq_len may contain garbage

722 + # physical page ids (see physical_to_logical_mapping). With batched

723 + # decode, max_seq_len is the batch max while shorter requests still

724 + # index all columns up to that max unless masked here.

725 + num_blocks = self.num_blocks_per_seq[self.doc_ids]

726 + past_seq = self.logical_block_ids[None, :] >= num_blocks[:, None]

727 + used_pages.masked_fill_(past_seq, 0)

728 +

729 custom_hint = self.block_sparsity_hint is not None

730 + use_rswa = self.rswa_window is not None and self.rswa_prefix_lens is not None

731 + needs_per_q_pruning = (

732 + self.causal or self.sliding_window or custom_hint or use_rswa

733 + )

734

668 - if self.sliding_window or custom_hint:

735 + if needs_per_q_pruning:

736 device = used_pages.device

737 assert self.doc_ids is not None

738 token_indices = torch.arange(

@@ -676,6 +743,12 @@ class FlexAttentionMetadata:

743 - self.query_start_loc[self.doc_ids]

744 + self.decode_offset[self.doc_ids]

745 )

746 + block_starts = self.logical_block_ids * self.block_size

747 + block_ends = block_starts + self.block_size

748 +

749 + if self.causal:

750 + future_blocks = block_starts[None, :] > logical_q_idx[:, None]

751 + used_pages.masked_fill_(future_blocks, 0)

752

753 if self.sliding_window:

754 assert self.sliding_window is not None

@@ -685,6 +758,23 @@ class FlexAttentionMetadata:

758 min_block_idx = min_kv_idx // self.block_size

759 sliding_mask = self.logical_block_ids >= min_block_idx[:, None]

760 used_pages.masked_fill_(~sliding_mask, 0)

761 + if use_rswa:

762 + # R-SWA keeps prefix KV globally visible and applies a sliding

763 + # window over generated tokens. Prune blocks that fall entirely

764 + # in the "hole" between prefix_len and the current window so

765 + # FlexAttention does not gather invalid paged-KV slots (this

766 + # mirrors uniform sliding-window block pruning above).

767 + assert self.rswa_prefix_lens is not None

768 + assert self.rswa_window is not None

769 + prefix_len = self.rswa_prefix_lens[self.doc_ids]

770 + min_kv_window = torch.maximum(

771 + prefix_len,

772 + logical_q_idx - (self.rswa_window - 1),

773 + )

774 + in_gap = (block_starts[None, :] >= prefix_len[:, None]) & (

775 + block_ends[None, :] <= min_kv_window[:, None]

776 + )

777 + used_pages.masked_fill_(in_gap, 0)

778 if custom_hint:

779 assert self.block_sparsity_hint is not None

780 q_block_idx = logical_q_idx // self.block_size

@@ -798,12 +888,36 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

888 self.max_num_query_groups = cdiv(max_num_batched_tokens, self.q_block_size)

889 max_num_pages_per_seq = cdiv(self.max_model_len, self.block_size)

890 self.max_num_kv_indices = self.q_block_size * max_num_pages_per_seq

891 + # R-SWA uses q_block_size=1 so block lists are not merged across requests

892 + # in a q-group (mixed-length batches otherwise gather foreign paged-KV).

893 + self.max_num_rswa_query_groups = max_num_batched_tokens

894 + # +1 sentinel column: the flex-attention kernel's get_offset_for_next_block

895 + # always prefetches kv_indices[q, kv_num_blocks] (one past the last valid

896 + # entry) to compute the jump offset for the next loop iteration. When

897 + # kv_num_blocks[q] == W (every page of the sequence is live), that prefetch

898 + # reads column W of the persistent buffer. Without the extra column this

899 + # would land on stale data from a previous step (the buffer is wider than W

900 + # but is never fully zeroed), producing an out-of-bounds K/V pointer and a

901 + # CUDA illegal memory access. Allocating W_max+1 columns and initialising

902 + # the whole buffer to -1 ensures the sentinel slot is always safe to read.

903 + self.max_num_rswa_kv_indices = max_num_pages_per_seq + 1

904 self.persistent_kv_num_blocks = torch.empty(

905 self.max_num_query_groups, dtype=torch.int32, device=device

906 )

907 + self.persistent_rswa_kv_num_blocks = torch.empty(

908 + self.max_num_rswa_query_groups, dtype=torch.int32, device=device

909 + )

910 self.persistent_offset_tensor = torch.empty(

911 max_num_seqs, dtype=torch.int32, device=device

912 )

913 + # Persistent buffer for R-SWA per-request prefix lengths so the device

914 + # address stays stable across steps (required for CUDA graph replay).

915 + self.rswa_window: int | None = self.model_config.rswa_window

916 + self.persistent_rswa_prefix_lens: torch.Tensor | None = None

917 + if self.rswa_window is not None:

918 + self.persistent_rswa_prefix_lens = torch.empty(

919 + max_num_seqs, dtype=torch.int32, device=device

920 + )

921 self.persistent_doc_ids = torch.empty(

922 max_num_batched_tokens, dtype=torch.int32, device=device

923 )

@@ -811,6 +925,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

925 # initialize later when we can access block_table

926 self.persistent_physical_to_logical = None

927 self.persistent_kv_indices = None

928 + self.persistent_rswa_kv_indices = None

929

930 self.custom_logical_mask_mod: _mask_mod_signature | None = None

931 if self._uses_full_cudagraphs():

@@ -936,6 +1051,26 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

1051 dtype=torch.int32,

1052 device=self.device,

1053 )

1054 + if self.persistent_rswa_kv_indices is None:

1055 + # Initialise to -1 so the +1 sentinel column (see max_num_rswa_kv_indices)

1056 + # is always a safe pad value for the flex kernel's prefetch.

1057 + self.persistent_rswa_kv_indices = torch.full(

1058 + (self.max_num_rswa_query_groups, self.max_num_rswa_kv_indices),

1059 + fill_value=-1,

1060 + dtype=torch.int32,

1061 + device=self.device,

1062 + )

1063 +

1064 + use_rswa = self.rswa_window is not None

1065 + q_block_size = 1 if use_rswa else self.q_block_size

1066 + persistent_kv_indices = (

1067 + self.persistent_rswa_kv_indices if use_rswa else self.persistent_kv_indices

1068 + )

1069 + persistent_kv_num_blocks = (

1070 + self.persistent_rswa_kv_num_blocks

1071 + if use_rswa

1072 + else self.persistent_kv_num_blocks

1073 + )

1074

1075 inverse_block_table = copy_to_persistent(

1076 self.persistent_physical_to_logical, inverse_block_table

@@ -944,6 +1079,13 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

1079 offset_tensor = common_attn_metadata.compute_num_computed_tokens()

1080 offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor)

1081

1082 + rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens

1083 + if use_rswa and rswa_prefix_lens is not None:

1084 + assert self.persistent_rswa_prefix_lens is not None

1085 + rswa_prefix_lens = copy_to_persistent(

1086 + self.persistent_rswa_prefix_lens, rswa_prefix_lens

1087 + )

1088 +

1089 uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec)

1090 logical_mask_mod = (

1091 bidirectional_mask_mod

@@ -986,12 +1128,14 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

1128 # attention block mask for encoder-only models, disable it temporarily.

1129 # see: https://github.com/vllm-project/vllm/pull/27329#issuecomment-3431484053

1130 direct_build=self.direct_build and uses_paged_kv,

989 - q_block_size=self.q_block_size,

1131 + q_block_size=q_block_size,

1132 kv_block_size=self.kv_block_size,

991 - persistent_kv_indices=self.persistent_kv_indices,

992 - persistent_kv_num_blocks=self.persistent_kv_num_blocks,

1133 + persistent_kv_indices=persistent_kv_indices,

1134 + persistent_kv_num_blocks=persistent_kv_num_blocks,

1135 persistent_doc_ids=self.persistent_doc_ids,

1136 mm_prefix_range=common_attn_metadata.mm_req_doc_ranges,

1137 + rswa_prefix_lens=rswa_prefix_lens,

1138 + rswa_window=self.rswa_window,

1139 )

1140

1141 # Pre-build block_mask so it is ready before CUDA graph capture.

vllm/v1/core/kv_cache_coordinator.py

+10 -2

@@ -329,7 +329,10 @@ class KVCacheCoordinator(ABC):

329 ]

330

331 def remove_skipped_blocks(

332 - self, request_id: str, total_computed_tokens: int

332 + self,

333 + request_id: str,

334 + total_computed_tokens: int,

335 + num_prompt_tokens: int | None = None,

336 ) -> None:

337 """

338 Remove the blocks that are no longer needed from `blocks` and replace

@@ -339,9 +342,14 @@ class KVCacheCoordinator(ABC):

342 request_id: The request ID.

343 total_computed_tokens: The total number of computed tokens, including

344 local computed tokens and external computed tokens.

345 + num_prompt_tokens: Optional prompt length. R-SWA managers use this to

346 + free gap blocks between the prefill tail and decode window; other

347 + manager types ignore it.

348 """

349 for manager in self.single_type_managers:

344 - manager.remove_skipped_blocks(request_id, total_computed_tokens)

350 + manager.remove_skipped_blocks(

351 + request_id, total_computed_tokens, num_prompt_tokens

352 + )

353

354 def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:

355 """

vllm/v1/core/kv_cache_manager.py

+11 -3

@@ -398,7 +398,9 @@ class KVCacheManager:

398 # Should call this function before allocating new blocks to reduce

399 # the number of evicted blocks.

400 self.coordinator.remove_skipped_blocks(

401 - request.request_id, total_computed_tokens

401 + request.request_id,

402 + total_computed_tokens,

403 + num_prompt_tokens=request.num_prompt_tokens,

404 )

405

406 num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(

@@ -468,7 +470,10 @@ class KVCacheManager:

470 self.coordinator.free(request.request_id)

471

472 def remove_skipped_blocks(

471 - self, request_id: str, total_computed_tokens: int

473 + self,

474 + request_id: str,

475 + total_computed_tokens: int,

476 + num_prompt_tokens: int | None = None,

477 ) -> None:

478 """Remove the blocks that are no longer needed from `blocks` and replace

479 the removed blocks with null_block.

@@ -477,8 +482,11 @@ class KVCacheManager:

482 request_id: The request ID.

483 total_computed_tokens: The total number of computed tokens, including

484 local computed tokens and external computed tokens.

485 + num_prompt_tokens: Optional prompt length for R-SWA gap eviction.

486 """

481 - self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens)

487 + self.coordinator.remove_skipped_blocks(

488 + request_id, total_computed_tokens, num_prompt_tokens

489 + )

490

491 def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:

492 """Pop the request's bookkeeping and return its blocks without

vllm/v1/core/sched/scheduler.py

@@ -2342,6 +2342,7 @@ class Scheduler(SchedulerInterface):

2342 self.kv_cache_manager.remove_skipped_blocks(

2343 request_id=request.request_id,

2344 total_computed_tokens=request.num_computed_tokens,

2345 + num_prompt_tokens=request.num_prompt_tokens,

2346 )

2347

2348 block_ids = self.kv_cache_manager.get_block_ids(request.request_id)

vllm/v1/core/single_type_kv_cache_manager.py

+103 -19

@@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import (

20 KVCacheSpec,

21 MambaSpec,

22 MLAAttentionSpec,

23 + RSWASpec,

24 SinkFullAttentionSpec,

25 SlidingWindowMLASpec,

26 SlidingWindowSpec,

@@ -476,8 +477,38 @@ class SingleTypeKVCacheManager(ABC):

477

478 raise NotImplementedError

479

480 + def _remove_blocks_in_range(

481 + self,

482 + request_id: str,

483 + first_block: int,

484 + last_block: int,

485 + ) -> None:

486 + """Free blocks in ``[first_block, last_block)`` and replace with null_block.

487 +

488 + Iterates backward so newly-evictable tail blocks are reached even after

489 + earlier blocks in the range were nulled in a prior call.

490 + """

491 + if request_id not in self.req_to_blocks:

492 + return

493 + if first_block >= last_block:

494 + return

495 + blocks = self.req_to_blocks[request_id]

496 + last_block = min(last_block, len(blocks))

497 +

498 + freed: list[KVCacheBlock] = []

499 + for i in range(last_block - 1, first_block - 1, -1):

500 + if blocks[i] == self._null_block:

501 + break

502 + freed.append(blocks[i])

503 + blocks[i] = self._null_block

504 + if freed:

505 + self.block_pool.free_blocks(freed)

506 +

507 def remove_skipped_blocks(

480 - self, request_id: str, total_computed_tokens: int

508 + self,

509 + request_id: str,

510 + total_computed_tokens: int,

511 + num_prompt_tokens: int | None = None,

512 ) -> None:

513 """

514 Remove and free the blocks that are no longer needed for attention computation.

@@ -490,7 +521,11 @@ class SingleTypeKVCacheManager(ABC):

521 request_id: The request ID.

522 total_computed_tokens: The total number of computed tokens, including

523 local computed tokens and external computed tokens.

524 + num_prompt_tokens: Optional prompt length for attention types (e.g.

525 + R-SWA) that evict a middle gap rather than a head prefix. Ignored

526 + by the default implementation.

527 """

528 + del num_prompt_tokens

529 # Remove the blocks that will be skipped during attention computation.

530 num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)

531 if num_skipped_tokens <= 0:

@@ -506,18 +541,7 @@ class SingleTypeKVCacheManager(ABC):

541 # range), so we must cap to the number of blocks that currently exist for

542 # this request.

543 num_skipped_blocks = min(num_skipped_blocks, len(blocks))

509 - removed_blocks: list[KVCacheBlock] = []

510 - # Because the block starts from index 0, the num_skipped_block-th block

511 - # corresponds to index num_skipped_blocks - 1.

512 - for i in range(num_skipped_blocks - 1, -1, -1):

513 - if blocks[i] == self._null_block:

514 - # If the block is already a null block, the blocks before it

515 - # should also have been set to null blocks by the previous calls

516 - # to this function.

517 - break

518 - removed_blocks.append(blocks[i])

519 - blocks[i] = self._null_block

520 - self.block_pool.free_blocks(removed_blocks)

544 + self._remove_blocks_in_range(request_id, 0, num_skipped_blocks)

545

546 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:

547 """

@@ -598,6 +622,50 @@ class FullAttentionManager(SingleTypeKVCacheManager):

622 return num_common_blocks

623

624

625 +class RSWAManager(FullAttentionManager):

626 + """KV cache manager for Reference Sliding Window Attention (R-SWA).

627 +

628 + When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees

629 + gap blocks between the prefill tail and the current decode window. This

630 + bounds per-request KV memory at O(prefix_len + rswa_window) instead of

631 + growing linearly with decode length.

632 + """

633 +

634 + def __init__(self, kv_cache_spec: RSWASpec, **kwargs) -> None:

635 + super().__init__(kv_cache_spec, **kwargs)

636 + self.rswa_window: int = kv_cache_spec.rswa_window

637 +

638 + def remove_skipped_blocks(

639 + self,

640 + request_id: str,

641 + total_computed_tokens: int,

642 + num_prompt_tokens: int | None = None,

643 + ) -> None:

644 + """Free gap blocks that are no longer needed for attention.

645 +

646 + Gap = blocks entirely within

647 + [ceil(prefix_len / block_size) * block_size,

648 + max(prefix_len, total_computed_tokens - rswa_window))

649 +

650 + Freed blocks are replaced with null_block in req_to_blocks so the

651 + block_table passed to FA4 is valid (null_block KV is all-zero;

652 + rswa_mask_mod marks gap positions as non-visible so FA4 skips them).

653 + """

654 + if num_prompt_tokens is None:

655 + super().remove_skipped_blocks(

656 + request_id, total_computed_tokens, num_prompt_tokens

657 + )

658 + return

659 +

660 + bs = self.block_size

661 + # First block fully after the prefill boundary.

662 + first_gap_block = cdiv(num_prompt_tokens, bs)

663 + # Decode window start position; blocks before this are evictable.

664 + window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window)

665 + last_gap_block = window_start // bs # exclusive upper bound

666 + self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block)

667 +

668 +

669 class SlidingWindowManager(SingleTypeKVCacheManager):

670 def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:

671 super().__init__(kv_cache_spec, **kwargs)

@@ -1072,7 +1140,12 @@ class MambaManager(SingleTypeKVCacheManager):

1140

1141 return mask

1142

1075 - def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:

1143 + def remove_skipped_blocks(

1144 + self,

1145 + request_id: str,

1146 + num_computed_tokens: int,

1147 + num_prompt_tokens: int | None = None,

1148 + ) -> None:

1149 assert isinstance(self.kv_cache_spec, MambaSpec)

1150

1151 # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain

@@ -1082,7 +1155,9 @@ class MambaManager(SingleTypeKVCacheManager):

1155 # that we might actually need.

1156 num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks)

1157

1085 - super().remove_skipped_blocks(request_id, num_computed_tokens)

1158 + super().remove_skipped_blocks(

1159 + request_id, num_computed_tokens, num_prompt_tokens

1160 + )

1161 if self.mamba_cache_mode == "align":

1162 # `last_state_block_idx` refers to the block index allocated two steps ago.

1163 # The block allocated in the previous step is used to copy Mamba states

@@ -1401,10 +1476,16 @@ def get_manager_for_kv_cache_spec(

1476 assert manager_class is not None, (

1477 f"No manager registered for KVCacheSpec {type(kv_cache_spec)}"

1478 )

1404 - # SlidingWindow / ChunkedLocalAttention managers recycle blocks across

1405 - # chunks; the runtime admission cap must match the recycling-aware bound

1406 - # the startup pool sizer uses (single source of truth: the spec method).

1407 - if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):

1479 + # SlidingWindow / ChunkedLocalAttention managers recycle blocks;

1480 + # the runtime admission cap must match the recycling-aware bound the

1481 + # startup pool sizer uses (single source of truth: the spec method).

1482 + # R-SWA also recycles gap blocks but peak physical KV still fits the

1483 + # full-attention bound (prefix + window <= max_model_len), so it inherits

1484 + # FullAttentionSpec sizing without a separate admission cap.

1485 + if isinstance(

1486 + kv_cache_spec,

1487 + (SlidingWindowSpec, ChunkedLocalAttentionSpec),

1488 + ):

1489 kwargs["max_admission_blocks_per_request"] = (

1490 kv_cache_spec.max_admission_blocks_per_request(

1491 max_num_batched_tokens=max_num_batched_tokens,

@@ -1457,6 +1538,9 @@ def register_all_kvcache_specs(vllm_config):

1538 KVCacheSpecRegistry.register(

1539 MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec

1540 )

1541 + KVCacheSpecRegistry.register(

1542 + RSWASpec, RSWAManager, uniform_type_base_spec=FullAttentionSpec

1543 + )

1544 # NOTE(Mengqing): HiddenStateCacheSpec won't take part in

1545 # grouping, thus the uniform_type_base_spec is just a

1546 # placeholder.

vllm/v1/kv_cache_interface.py

+40

@@ -437,6 +437,46 @@ class HiddenStateCacheSpec(MLAAttentionSpec):

437 pass

438

439

440 +@dataclass(frozen=True, kw_only=True)

441 +class RSWASpec(FullAttentionSpec):

442 + """KV cache spec for Reference Sliding Window Attention (R-SWA).

443 +

444 + Prefill (image + text prompt) tokens are always globally visible.

445 + Only the last ``rswa_window`` generated tokens are kept in the KV cache;

446 + gap blocks (between the prefill tail and the current decode window) are

447 + evicted during each decode step to bound memory at

448 + O(prefix_blocks + window_blocks).

449 + """

450 +

451 + rswa_window: int

452 +

453 + @classmethod

454 + def merge(cls, specs: list[RSWASpec]) -> RSWASpec:

455 + assert all(isinstance(spec, RSWASpec) for spec in specs), (

456 + "All attention layers in the same KV cache group must be RSWASpec."

457 + )

458 + rswa_windows = {spec.rswa_window for spec in specs}

459 + assert len(rswa_windows) == 1, (

460 + f"All R-SWA layers must share the same rswa_window, got {rswa_windows}"

461 + )

462 + # Delegate common field merging to the parent, then reattach rswa_window.

463 + base = FullAttentionSpec.merge(specs) # type: ignore[arg-type]

464 + return cls(

465 + block_size=base.block_size,

466 + num_kv_heads=base.num_kv_heads,

467 + head_size=base.head_size,

468 + head_size_v=base.head_size_v,

469 + dtype=base.dtype,

470 + kv_quant_mode=base.kv_quant_mode,

471 + page_size_padded=base.page_size_padded,

472 + indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,

473 + sliding_window=base.sliding_window,

474 + attention_chunk_size=base.attention_chunk_size,

475 + non_causal=base.non_causal,

476 + rswa_window=rswa_windows.pop(),

477 + )

478 +

479 +

480 @dataclass(frozen=True, kw_only=True)

481 class ChunkedLocalAttentionSpec(AttentionSpec):

482 attention_chunk_size: int

vllm/v1/worker/gpu/attn_utils.py

@@ -469,6 +469,7 @@ def build_attn_metadata(

469 model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,

470 for_cudagraph_capture: bool = False,

471 causal: bool = True,

472 + rswa_prefix_lens: torch.Tensor | None = None,

473 ) -> dict[str, Any]:

474 seq_lens = seq_lens[:num_reqs]

475 if dcp_local_seq_lens is not None:

@@ -501,6 +502,7 @@ def build_attn_metadata(

502 causal=causal,

503 dcp_local_seq_lens=dcp_local_seq_lens,

504 positions=positions,

505 + rswa_prefix_lens=rswa_prefix_lens,

506 **common_attn_metadata_extra_kwargs,

507 )

508

vllm/v1/worker/gpu/input_batch.py

@@ -96,6 +96,9 @@ class InputBatch:

96 # Whether any requests in batch use structured output.

97 has_structured_output_reqs: bool

99 + # [num_reqs_after_padding] per-request prompt length for R-SWA (optional).

100 + rswa_prefix_lens: torch.Tensor | None = None

101 +

102 @classmethod

103 def make_dummy(

104 cls,

vllm/v1/worker/gpu/model_runner.py

+17

@@ -223,6 +223,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

223 max_num_tokens=self.max_num_tokens,

224 device=self.device,

225 )

226 + # R-SWA: persistent GPU buffer for per-request prefix lengths (CUDA-graph safe).

227 + self.rswa_prefix_lens_buffer: torch.Tensor | None = None

228 + if self.model_config.rswa_window is not None:

229 + self.rswa_prefix_lens_buffer = torch.zeros(

230 + self.max_num_reqs, dtype=torch.int32, device=self.device

231 + )

232

233 if self.use_pp:

234 self.pp_handler = PPHandler(

@@ -985,6 +991,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

991 if self.use_pp:

992 # max_seq_len is only consumed by the PP `compute_need_sampled_mask`

993 max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]

994 +

995 + rswa_prefix_lens = None

996 + if self.rswa_prefix_lens_buffer is not None:

997 + rswa_prefix_lens = self.rswa_prefix_lens_buffer[:num_reqs_padded]

998 + rswa_prefix_lens[:num_reqs] = self.req_states.prompt_len.gpu[

999 + idx_mapping[:num_reqs]

1000 + ]

1001 + if num_reqs_padded > num_reqs:

1002 + rswa_prefix_lens[num_reqs:].zero_()

1003 +

1004 return InputBatch(

1005 req_ids=req_ids,

1006 num_reqs=num_reqs,

@@ -1015,6 +1031,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1031 cu_num_logits=cu_num_logits,

1032 cu_num_logits_np=cu_num_logits_np,

1033 has_structured_output_reqs=scheduler_output.has_structured_output_requests,

1034 + rswa_prefix_lens=rswa_prefix_lens,

1035 )

1036

1037 def prepare_attn(

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

@@ -168,5 +168,6 @@ class DefaultModelState(ModelState):

168 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

169 positions=input_batch.positions,

170 for_cudagraph_capture=for_capture,

171 + rswa_prefix_lens=input_batch.rswa_prefix_lens,

172 )

173 return attn_metadata

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

@@ -146,6 +146,7 @@ class EncoderDecoderModelState(ModelState):

146 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

147 model_specific_attn_metadata=enc_dec_attn_metadata,

148 for_cudagraph_capture=for_capture,

149 + rswa_prefix_lens=input_batch.rswa_prefix_lens,

150 )

151 return attn_metadata

152

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

@@ -141,6 +141,7 @@ class MambaHybridModelState(DefaultModelState):

141 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

142 model_specific_attn_metadata=mamba_attn_metadata,

143 for_cudagraph_capture=for_capture,

144 + rswa_prefix_lens=input_batch.rswa_prefix_lens,

145 )

146

147 def postprocess_state(

vllm/v1/worker/gpu_model_runner.py

@@ -2335,6 +2335,13 @@ class GPUModelRunner(

2335 req_idx = self.input_batch.req_id_to_index[req_id]

2336 req_doc_ranges[req_idx] = image_doc_ranges

2337

2338 + # Reference Sliding Window Attention (R-SWA): pass per-request prompt

2339 + # lengths so the attention backend can keep the prefix globally visible.

2340 + # The backend owns the persistent CUDA-graph-safe GPU buffer.

2341 + rswa_prefix_lens = None

2342 + if self.model_config.rswa_window is not None:

2343 + rswa_prefix_lens = num_prompt_tokens_cpu

2344 +

2345 cm_base = CommonAttentionMetadata(

2346 query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],

2347 query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],

@@ -2352,6 +2359,7 @@ class GPUModelRunner(

2359 is_prefilling=is_prefilling,

2360 positions=self.positions[:num_tokens_padded],

2361 mm_req_doc_ranges=req_doc_ranges,

2362 + rswa_prefix_lens=rswa_prefix_lens,

2363 )

2364

2365 if self.dcp_world_size > 1: