@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 17dc9c7fc9
Fork 0

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):