[ModelRunner V2] Simplify recent UnlimitedOCR-related changes (#46975)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 29, 2026 at 17:56 UTC
8ad4a01825ef941e785b2bc305ac7a6b5ca9c530
5 files changed
+11
-21
vllm/v1/worker/gpu/input_batch.py
+3
-2
@@ -96,8 +96,8 @@ class InputBatch:
96
# Whether any requests in batch use structured output.
97
has_structured_output_reqs: bool
98
99
- # [num_reqs_after_padding] per-request prompt length for R-SWA (optional).
100
- rswa_prefix_lens: torch.Tensor | None = None
99
+ # [num_reqs] per-request prompt length, only populated for R-SWA.
100
+ prompt_lens: torch.Tensor | None
101
102
@classmethod
103
def make_dummy(
@@ -178,6 +178,7 @@ class InputBatch:
178
cu_num_logits=cu_num_logits,
179
cu_num_logits_np=cu_num_logits_np,
180
has_structured_output_reqs=False,
181
+ prompt_lens=None,
182
)
183
184
vllm/v1/worker/gpu/model_runner.py
+5
-16
@@ -223,13 +223,6 @@ 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
-
226
if self.use_pp:
227
self.pp_handler = PPHandler(
228
max_num_reqs=self.max_num_reqs,
@@ -992,14 +985,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
985
# max_seq_len is only consumed by the PP `compute_need_sampled_mask`
986
max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
987
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_()
988
+ prompt_lens = None
989
+ if self.model_config.rswa_window is not None:
990
+ # prompt_lens is only used in R-SWA case.
991
+ prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]
992
993
return InputBatch(
994
req_ids=req_ids,
@@ -1031,7 +1020,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1020
cu_num_logits=cu_num_logits,
1021
cu_num_logits_np=cu_num_logits_np,
1022
has_structured_output_reqs=scheduler_output.has_structured_output_requests,
1034
- rswa_prefix_lens=rswa_prefix_lens,
1023
+ prompt_lens=prompt_lens,
1024
)
1025
1026
def prepare_attn(
vllm/v1/worker/gpu/model_states/default.py
+1
-1
@@ -168,6 +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,
171
+ rswa_prefix_lens=input_batch.prompt_lens,
172
)
173
return attn_metadata
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+1
-1
@@ -146,7 +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,
149
+ rswa_prefix_lens=input_batch.prompt_lens,
150
)
151
return attn_metadata
152
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+1
-1
@@ -141,7 +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,
144
+ rswa_prefix_lens=input_batch.prompt_lens,
145
)
146
147
def postprocess_state(