A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
[CI] Bump `mypy` version (#34950)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 4, 2026 at 20:55 UTC
17dc9c7fc94534e542b6849192ed382c122d2d08
13 files changed
+90
-61
.pre-commit-config.yaml
+1
-1
@@ -55,7 +55,7 @@ repos:
55
language: python
56
types_or: [python, pyi]
57
require_serial: true
58
- additional_dependencies: [mypy==1.11.1, regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
58
+ additional_dependencies: ["mypy[faster-cache]==1.15.0", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
59
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
60
name: Run mypy for Python 3.10
61
entry: python tools/pre_commit/mypy.py 1 "3.10"
tests/kernels/core/test_pos_encoding.py
+2
-5
@@ -94,12 +94,9 @@ def test_rotary_embedding(
94
95
positions = torch.randint(0, max_position, (batch_size, seq_len))
96
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
97
- query = torch.randn(query_shape, dtype=dtype)
98
- key = torch.randn_like(query) if use_key else None
99
-
97
# slice tensor if required, noop otherwise
101
- query = query[..., :head_size]
102
- key = key[..., :head_size] if use_key else None
98
+ query = torch.randn(query_shape, dtype=dtype)[..., :head_size]
99
+ key = torch.randn_like(query)[..., :head_size] if use_key else None
100
101
# NOTE(woosuk): The reference implementation should be executed first
102
# because the custom kernel is in-place.
tests/kernels/core/test_rotary_embedding.py
+2
-2
@@ -62,7 +62,7 @@ def test_rotary_embedding_opcheck(
62
)
63
key = torch.randn_like(query) if use_key else None
64
query = query[..., :head_size]
65
- key = key[..., :head_size] if use_key else None
65
+ key = key[..., :head_size] if key is not None else None
66
67
rotary_embedding_opcheck(rot, positions, query, key)
68
@@ -73,5 +73,5 @@ def test_rotary_embedding_opcheck(
73
rot,
74
positions,
75
query.flatten(start_dim=-2),
76
- key.flatten(start_dim=-2) if use_key else None,
76
+ key.flatten(start_dim=-2) if key is not None else None,
77
)
tests/kernels/mamba/test_mamba_ssm.py
+7
-7
@@ -298,13 +298,13 @@ def test_selective_scan(
298
C = torch.randn(C_shape, device=device, dtype=wtype if not is_variable_C else itype)
299
C_ref = C.clone()
300
D = torch.randn(dim, device=device, dtype=torch.float32) if has_D else None
301
- D_ref = D.clone()
301
+ D_ref = D.clone() if D is not None else None
302
z = (
303
torch.randn(batch_size, dim, seqlen, device=device, dtype=itype)
304
if has_z
305
else None
306
)
307
- z_ref = z.clone() if has_z else None
307
+ z_ref = z.clone() if z is not None else None
308
delta_bias = (
309
(0.5 * torch.rand(dim, device=device, dtype=torch.float32))
310
if has_delta_bias
@@ -493,7 +493,7 @@ def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
493
B[idx : idx + 1],
494
C[idx : idx + 1],
495
D=D,
496
- z=z[idx : idx + 1] if has_z else None,
496
+ z=z[idx : idx + 1] if z is not None else None,
497
dt_bias=dt_bias,
498
dt_softplus=True,
499
)
@@ -578,7 +578,7 @@ def test_selective_scan_varlen(
578
C = torch.randn(C_shape, device=device, dtype=wtype if not is_variable_C else itype)
579
C_ref = C.clone()
580
D = torch.randn(dim, device=device, dtype=torch.float32) if has_D else None
581
- D_ref = D.clone()
581
+ D_ref = D.clone() if D is not None else None
582
z = torch.randn(dim, seqlen, device=device, dtype=itype)
583
z_ref = z.clone()
584
delta_bias = (
@@ -750,7 +750,7 @@ def test_selective_state_update_with_batch_indices(
750
B[:batch_size],
751
C[:batch_size],
752
D=D,
753
- z=z[:batch_size],
753
+ z=z[:batch_size] if z is not None else None,
754
dt_bias=dt_bias,
755
dt_softplus=True,
756
)
@@ -934,7 +934,7 @@ def test_selective_state_update_with_num_accepted_tokens(
934
B[global_idx : global_idx + 1],
935
C[global_idx : global_idx + 1],
936
D=D,
937
- z=z[global_idx : global_idx + 1] if has_z else None,
937
+ z=z[global_idx : global_idx + 1] if z is not None else None,
938
dt_bias=dt_bias,
939
dt_softplus=True,
940
)
@@ -1061,7 +1061,7 @@ def test_selective_state_update_varlen_with_num_accepted(
1061
B[global_idx : global_idx + 1],
1062
C[global_idx : global_idx + 1],
1063
D=D,
1064
- z=z[global_idx : global_idx + 1] if has_z else None,
1064
+ z=z[global_idx : global_idx + 1] if z is not None else None,
1065
dt_bias=dt_bias,
1066
dt_softplus=True,
1067
)
tests/kernels/quantization/test_fp8_quant.py
+3
-3
@@ -57,11 +57,11 @@ def opcheck_fp8_quant(
57
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
58
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
59
@pytest.mark.parametrize("dtype", DTYPES)
60
-@pytest.mark.parametrize("scale_ub", SCALE_UBS)
60
+@pytest.mark.parametrize("do_scale_ub", SCALE_UBS)
61
@pytest.mark.parametrize("seed", SEEDS)
62
@torch.inference_mode()
63
def test_dynamic_per_token_fp8_quant(
64
- num_tokens: int, hidden_size: int, dtype: torch.dtype, scale_ub: bool, seed: int
64
+ num_tokens: int, hidden_size: int, dtype: torch.dtype, do_scale_ub: bool, seed: int
65
) -> None:
66
set_random_seed(seed)
67
@@ -70,7 +70,7 @@ def test_dynamic_per_token_fp8_quant(
70
) # avoid nans
71
72
scale_ub = (
73
- torch.mean(x).to(dtype=torch.float32, device="cuda") if scale_ub else None
73
+ torch.mean(x).to(dtype=torch.float32, device="cuda") if do_scale_ub else None
74
)
75
ref_out, ref_scales = ref_dynamic_per_token_quant(x, FP8_DTYPE, scale_ub)
76
ops_out, ops_scales = ops.scaled_fp8_quant(
vllm/config/parallel.py
+13
-3
@@ -3,11 +3,11 @@
3
4
import os
5
from collections.abc import Callable
6
-from typing import TYPE_CHECKING, Any, Literal
6
+from typing import TYPE_CHECKING, Any, Literal, overload
7
8
import torch
9
from pydantic import Field, field_validator, model_validator
10
-from torch.distributed import ProcessGroup, ReduceOp
10
+from torch.distributed import ProcessGroup, ReduceOp, Store
11
from typing_extensions import Self
12
13
import vllm.envs as envs
@@ -507,7 +507,17 @@ class ParallelConfig:
507
def get_next_stateless_eplb_group_port(self) -> list[int]:
508
return self._stateless_eplb_group_port_list.pop()
509
510
- def stateless_init_dp_group(self, return_store: bool = False) -> ProcessGroup:
510
+ @overload
511
+ def stateless_init_dp_group(
512
+ self, return_store: Literal[False] = ...
513
+ ) -> ProcessGroup: ...
514
+ @overload
515
+ def stateless_init_dp_group(
516
+ self, return_store: Literal[True] = ...
517
+ ) -> tuple[ProcessGroup, Store]: ...
518
+ def stateless_init_dp_group(
519
+ self, return_store: bool = False
520
+ ) -> ProcessGroup | tuple[ProcessGroup, Store]:
521
# NOTE: In high-concurrency scenarios multiple processes
522
# can pick the same (currently free) port through a race
523
# condition when calling `get_open_port()`. When the first
vllm/distributed/elastic_ep/elastic_state.py
+25
-11
@@ -4,7 +4,7 @@ import enum
4
import time
5
import weakref
6
from datetime import timedelta
7
-from typing import TYPE_CHECKING, Literal
7
+from typing import TYPE_CHECKING, Literal, TypeAlias
8
9
import torch.distributed
10
@@ -61,6 +61,14 @@ class ScaleDownRemovingEngineState(enum.IntEnum):
61
COMPLETE = 2
62
63
64
+EngineState: TypeAlias = (
65
+ ScaleUpExistingEngineState
66
+ | ScaleUpNewEngineState
67
+ | ScaleDownRemainingEngineState
68
+ | ScaleDownRemovingEngineState
69
+)
70
+
71
+
72
class _BarrierTimeoutError(RuntimeError):
73
"""
74
Exception raised for timeout
@@ -87,14 +95,13 @@ class ElasticEPScalingState:
95
self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None
96
self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None
97
self.new_parallel_config: ParallelConfig = new_parallel_config
90
- self.new_dp_group: torch.distributed.ProcessGroup | None = (
91
- self.engine_core.dp_group if worker_type == "new" else None
92
- )
98
+ self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None
99
self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None
100
self.worker_type = worker_type
101
self.scale_type = scale_type
102
self.reconfig_request = reconfig_request
103
104
+ self.state: EngineState
105
if scale_type == "scale_up":
106
self.state = (
107
ScaleUpNewEngineState.PREPARE
@@ -182,9 +189,9 @@ class ElasticEPScalingState:
189
engine step, and will synchronize with the other EngineCores in the
190
next step with a barrier without timeout.
191
"""
185
- dp_store = self.new_dp_store if use_new_group else self.old_dp_store
192
dp_group = self.new_dp_group if use_new_group else self.old_dp_group
187
- assert dp_group is not None
193
+ dp_store = self.new_dp_store if use_new_group else self.old_dp_store
194
+ assert dp_group is not None and dp_store is not None
195
196
group_rank = dp_group.rank()
197
group_size = dp_group.size()
@@ -212,6 +219,7 @@ class ElasticEPScalingState:
219
220
def _progress_existing_engine(self) -> bool:
221
state = self.state
222
+ assert self.old_dp_group is not None and self.old_dp_store is not None
223
224
if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT:
225
return False
@@ -265,11 +273,12 @@ class ElasticEPScalingState:
273
elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE:
274
self._switch_and_prepare()
275
self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE
276
+ assert self.new_dp_store is not None
277
self.new_dp_store.add("eep_barrier_engine_count", 1)
278
return True
279
280
elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE:
272
- assert self.new_dp_group is not None
281
+ assert self.new_dp_group is not None and self.new_dp_store is not None
282
if (
283
int(self.new_dp_store.get("eep_barrier_engine_count"))
284
< self.new_dp_group.size()
@@ -292,7 +301,7 @@ class ElasticEPScalingState:
301
302
def _progress_new_engine(self) -> bool:
303
state = self.state
295
- assert self.new_dp_group is not None
304
+ assert self.new_dp_group is not None and self.new_dp_store is not None
305
306
if state == ScaleUpNewEngineState.PREPARE:
307
tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")
@@ -330,6 +339,7 @@ class ElasticEPScalingState:
339
340
def _progress_remaining_engine(self) -> bool:
341
state = self.state
342
+ assert self.old_dp_group is not None and self.old_dp_store is not None
343
344
if state == ScaleDownRemainingEngineState.PREPARE:
345
self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE
@@ -369,6 +379,7 @@ class ElasticEPScalingState:
379
380
def _progress_removing_engine(self) -> bool:
381
state = self.state
382
+ assert self.old_dp_group is not None and self.old_dp_store is not None
383
384
if state == ScaleDownRemovingEngineState.PREPARE:
385
self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE
@@ -401,6 +412,7 @@ class ElasticEPScalingState:
412
413
def handle_notification(self, notification_type: EEPNotificationType):
414
assert self.worker_type != "new"
415
+ assert self.old_dp_store is not None
416
if (
417
notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY
418
and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
@@ -429,6 +441,7 @@ class ElasticEPScalingState:
441
)
442
443
def _create_standby_groups(self):
444
+ assert self.old_dp_group is not None
445
self.new_dp_group, self.new_dp_store = (
446
self.new_parallel_config.stateless_init_dp_group(return_store=True)
447
)
@@ -439,7 +452,7 @@ class ElasticEPScalingState:
452
logger.info("[Elastic EP] Created standby communication groups")
453
454
def _transfer_weights(self):
442
- assert self.reconfig_request is not None
455
+ assert self.reconfig_request is not None and self.old_dp_group is not None
456
old_dp_size = self.old_dp_group.size()
457
new_dp_size = self.reconfig_request.new_data_parallel_size
458
@@ -450,6 +463,7 @@ class ElasticEPScalingState:
463
logger.info("[Elastic EP] Transferred weights to new workers")
464
465
def _transfer_expert_mapping(self):
466
+ assert self.old_dp_group is not None
467
self.model_executor.collective_rpc(
468
"elastic_ep_execute", args=("broadcast_expert_mapping",)
469
)
@@ -458,7 +472,7 @@ class ElasticEPScalingState:
472
473
def _sync_kv_cache_memory_size(self):
474
assert self.engine_core.available_gpu_memory_for_kv_cache > 0
461
- assert self.new_dp_group is not None
475
+ assert self.new_dp_group is not None and self.old_dp_group is not None
476
ParallelConfig.sync_kv_cache_memory_size(
477
self.new_dp_group,
478
self.engine_core.available_gpu_memory_for_kv_cache,
@@ -507,7 +521,7 @@ class ElasticEPScalingState:
521
logger.info("[Elastic EP] EPLB reshuffle completed")
522
523
def _eplb_reshuffle_before_scale_down(self):
510
- assert self.reconfig_request is not None
524
+ assert self.reconfig_request is not None and self.old_dp_group is not None
525
self.model_executor.collective_rpc(
526
"elastic_ep_execute",
527
args=(
vllm/distributed/kv_transfer/kv_connector/utils.py
+1
@@ -336,6 +336,7 @@ class TpKVTopology:
336
self._cross_layers_blocks = (
337
len(self.tensor_shape) == len(kv_cache_shape) + 1
338
)
339
+ self.tensor_shape: torch.Size
340
341
if self._cross_layers_blocks:
342
logger.debug("Using cross-layer KV cache")
vllm/v1/attention/backends/flashinfer.py
+2
@@ -972,6 +972,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
972
973
# Early-out for cascade attention
974
if use_cascade:
975
+ assert num_blocks_np is not None
976
# Grab the blocks of the shared prefix from the first request.
977
num_common_kv_blocks = common_prefix_len // page_size
978
@@ -1117,6 +1118,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1118
max_seq_len=max_seq_len,
1119
)
1120
else:
1121
+ assert seq_lens_cpu is not None
1122
pure_decode = num_prefills == 0
1123
use_cudagraph = (
1124
self.enable_cuda_graph
vllm/v1/attention/backends/gdn_attn.py
+12
-11
@@ -88,14 +88,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
88
self.num_spec: int = self.speculative_config.num_speculative_tokens
89
else:
90
self.num_spec = 0
91
- self.use_spec_decode = self.num_spec > 0
91
+ self.use_spec_decode: bool = self.num_spec > 0
92
self._init_reorder_batch_threshold(1, self.use_spec_decode)
93
94
- self.use_full_cuda_graph = (
94
+ self.use_full_cuda_graph: bool = (
95
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
96
)
97
98
- self.decode_cudagraph_max_bs = (
98
+ self.decode_cudagraph_max_bs: int = (
99
self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)
100
)
101
if self.compilation_config.max_cudagraph_capture_size is not None:
@@ -104,42 +104,42 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
104
self.compilation_config.max_cudagraph_capture_size,
105
)
106
107
- self.spec_state_indices_tensor = torch.empty(
107
+ self.spec_state_indices_tensor: torch.Tensor = torch.empty(
108
(self.decode_cudagraph_max_bs, self.num_spec + 1),
109
dtype=torch.int32,
110
device=device,
111
)
112
- self.non_spec_state_indices_tensor = torch.empty(
112
+ self.non_spec_state_indices_tensor: torch.Tensor = torch.empty(
113
(self.decode_cudagraph_max_bs,),
114
dtype=torch.int32,
115
device=device,
116
)
117
- self.spec_sequence_masks = torch.empty(
117
+ self.spec_sequence_masks: torch.Tensor = torch.empty(
118
(self.decode_cudagraph_max_bs,),
119
dtype=torch.bool,
120
device=device,
121
)
122
- self.spec_token_indx = torch.empty(
122
+ self.spec_token_indx: torch.Tensor = torch.empty(
123
(self.decode_cudagraph_max_bs * (self.num_spec + 1),),
124
dtype=torch.int32,
125
device=device,
126
)
127
- self.non_spec_token_indx = torch.empty(
127
+ self.non_spec_token_indx: torch.Tensor = torch.empty(
128
(self.decode_cudagraph_max_bs * (self.num_spec + 1),),
129
dtype=torch.int32,
130
device=device,
131
)
132
- self.spec_query_start_loc = torch.empty(
132
+ self.spec_query_start_loc: torch.Tensor = torch.empty(
133
(self.decode_cudagraph_max_bs + 1,),
134
dtype=torch.int32,
135
device=device,
136
)
137
- self.non_spec_query_start_loc = torch.empty(
137
+ self.non_spec_query_start_loc: torch.Tensor = torch.empty(
138
(self.decode_cudagraph_max_bs + 1,),
139
dtype=torch.int32,
140
device=device,
141
)
142
- self.num_accepted_tokens = torch.empty(
142
+ self.num_accepted_tokens: torch.Tensor = torch.empty(
143
(self.decode_cudagraph_max_bs,),
144
dtype=torch.int32,
145
device=device,
@@ -322,6 +322,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
322
and num_spec_decodes <= self.decode_cudagraph_max_bs
323
and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
324
):
325
+ assert spec_sequence_masks is not None
326
self.spec_state_indices_tensor[:num_spec_decodes].copy_(
327
spec_state_indices_tensor, non_blocking=True
328
)
vllm/v1/attention/backends/mamba_attn.py
+6
-6
@@ -98,8 +98,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
98
self.use_spec_decode = self.num_spec_tokens > 0
99
100
assert isinstance(kv_cache_spec, MambaSpec)
101
- self.compilation_config = vllm_config.compilation_config
102
- self.decode_cudagraph_max_bs = self.vllm_config.scheduler_config.max_num_seqs
101
+ scheduler_config = vllm_config.scheduler_config
102
+ self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs
103
if self.compilation_config.max_cudagraph_capture_size is not None:
104
self.decode_cudagraph_max_bs = min(
105
self.decode_cudagraph_max_bs,
@@ -114,7 +114,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
114
# Speculative decoding not supported with prefix caching,
115
# so keep shape consistent with prefill buffer
116
# TODO: reduce this size as needed for decode-only cudagraph capture
117
- self.state_indices_tensor_d = torch.empty(
117
+ self.state_indices_tensor_d: torch.Tensor = torch.empty(
118
(
119
self.decode_cudagraph_max_bs,
120
max_num_blocks,
@@ -122,12 +122,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
122
dtype=torch.int32,
123
device=device,
124
)
125
- self.block_idx_last_scheduled_token = torch.empty(
125
+ self.block_idx_last_scheduled_token: torch.Tensor = torch.empty(
126
(self.decode_cudagraph_max_bs,),
127
dtype=torch.int32,
128
device=device,
129
)
130
- self.block_idx_last_computed_token = torch.empty(
130
+ self.block_idx_last_computed_token: torch.Tensor = torch.empty(
131
(self.decode_cudagraph_max_bs,),
132
dtype=torch.int32,
133
device=device,
@@ -142,7 +142,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
142
# For speculative decoding, we need to store the following buffers
143
# for CUDA graph capture during decode
144
if self.num_spec_tokens > 0:
145
- self.decode_num_accepted_tokens = torch.empty(
145
+ self.decode_num_accepted_tokens: torch.Tensor = torch.empty(
146
(self.decode_cudagraph_max_bs,),
147
dtype=torch.int32,
148
device=device,
vllm/v1/engine/core.py
+6
-6
@@ -1539,18 +1539,18 @@ class DPEngineCoreProc(EngineCoreProc):
1539
1540
def _init_data_parallel(self, vllm_config: VllmConfig):
1541
# Configure GPUs and stateless process group for data parallel.
1542
- dp_rank = vllm_config.parallel_config.data_parallel_rank
1543
- dp_size = vllm_config.parallel_config.data_parallel_size
1544
- local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local
1542
+ parallel_config = vllm_config.parallel_config
1543
+ dp_rank = parallel_config.data_parallel_rank
1544
+ dp_size = parallel_config.data_parallel_size
1545
+ local_dp_rank = parallel_config.data_parallel_rank_local
1546
1547
assert dp_size > 1
1548
assert local_dp_rank is not None
1549
assert 0 <= local_dp_rank <= dp_rank < dp_size
1550
1551
self.dp_rank = dp_rank
1551
- self.dp_group, self.dp_store = (
1552
- vllm_config.parallel_config.stateless_init_dp_group(return_store=True)
1553
- )
1552
+ dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True)
1553
+ self.dp_group, self.dp_store = dp_group, dp_store
1554
1555
def shutdown(self):
1556
super().shutdown()
vllm/v1/sample/logits_processor/__init__.py
+10
-6
@@ -309,12 +309,16 @@ class AdapterLogitsProcessor(LogitsProcessor):
309
310
"""
311
if req_lp := self.new_req_logits_processor(params):
312
- args = (
313
- [prompt_ids, output_ids]
314
- if (len(inspect.signature(req_lp).parameters) == 3)
315
- else [output_ids]
316
- )
317
- return partial(req_lp, *args) # type: ignore[misc]
312
+ if len(inspect.signature(req_lp).parameters) == 3:
313
+ if prompt_ids is None:
314
+ raise ValueError(
315
+ "Prompt token ids are required for this "
316
+ "logits processor but were not provided."
317
+ )
318
+ args = [prompt_ids, output_ids]
319
+ else:
320
+ args = [output_ids]
321
+ return partial(req_lp, *args)
322
return None
323
324
def update_state(self, batch_update: BatchUpdate | None):