From 4f65a6466656cda958ab650198b7237062e9904f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 10 Nov 2025 00:42:50 +0800 Subject: [PATCH] Refactor / Unify event loop across PD-Disagg, Overlap, DP-Attn cases (#12839) Co-authored-by: cctry <17473714+cctry@users.noreply.github.com> --- python/sglang/srt/disaggregation/decode.py | 142 ++++-------------- .../decode_schedule_batch_mixin.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 1 + python/sglang/srt/managers/scheduler.py | 8 +- .../srt/managers/scheduler_dp_attn_mixin.py | 65 +++++--- .../scheduler_output_processor_mixin.py | 22 ++- .../scheduler_runtime_checker_mixin.py | 11 ++ .../srt/model_executor/forward_batch_info.py | 7 + python/sglang/srt/two_batch_overlap.py | 6 +- test/srt/run_suite.py | 2 +- test/srt/test_disaggregation_dp_attention.py | 20 ++- 11 files changed, 142 insertions(+), 144 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 328a22adf..e9b9161a7 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -48,6 +48,7 @@ from sglang.srt.disaggregation.utils import ( ) from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, ScheduleBatch +from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ( @@ -58,11 +59,7 @@ from sglang.srt.mem_cache.memory_pool import ( ReqToTokenPool, SWAKVPool, ) -from sglang.srt.tracing.trace import ( - trace_event_batch, - trace_slice_batch, - trace_slice_end, -) +from sglang.srt.tracing.trace import trace_event_batch, trace_slice_end from sglang.srt.utils import get_int_env_var from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -587,11 +584,11 @@ class DecodePreallocQueue: need_space_for_single_req, ) - # Note: if the last fake extend just finishes, and we enter `pop_preallocated` immediately in the next iteration + # Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration # the extend batch is not in any queue, so we need to explicitly add the tokens slots here if ( self.scheduler.last_batch - and self.scheduler.last_batch.forward_mode.is_extend() + and self.scheduler.last_batch.forward_mode.is_prebuilt_extend() ): allocatable_tokens -= self.num_reserved_decode_tokens * len( self.scheduler.last_batch.reqs @@ -773,29 +770,12 @@ class DecodeTransferQueue: indices_to_remove.add(i) decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter() - - decode_req.req.check_finished() - if decode_req.req.finished(): - # finish immediately - decode_req.req.time_stats.forward_entry_time = ( - decode_req.req.time_stats.completion_time - ) = time.perf_counter() - self.scheduler.stream_output( - [decode_req.req], decode_req.req.return_logprob - ) - self.tree_cache.cache_finished_req(decode_req.req) - trace_slice_end( - RequestStage.DECODE_QUICK_FINISH, - decode_req.req.rid, - thread_finish_flag=True, - ) - else: - transferred_reqs.append(decode_req.req) - trace_slice_end( - RequestStage.DECODE_TRANSFERRED, - decode_req.req.rid, - auto_next_anon=True, - ) + transferred_reqs.append(decode_req.req) + trace_slice_end( + RequestStage.DECODE_TRANSFERRED, + decode_req.req.rid, + auto_next_anon=True, + ) elif poll in [ KVPoll.Bootstrapping, @@ -835,31 +815,9 @@ class SchedulerDisaggregationDecodeMixin: if batch: # Generate fake extend output. - if batch.forward_mode.is_extend(): - # Note: Logprobs should be handled on the prefill engine. - self.stream_output( - batch.reqs, any(req.return_logprob for req in batch.reqs) - ) - trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) - if self.require_mlp_sync: - self._prepare_idle_batch_and_run() - else: - if self.require_mlp_sync: - self.prepare_mlp_sync_batch(batch) - result = self.run_batch(batch) - self.process_batch_result(batch, result) - elif self.require_mlp_sync: - batch, _ = self._prepare_idle_batch_and_run() - - queue_size = ( - len(self.waiting_queue) - + len(self.disagg_decode_transfer_queue.queue) - + len(self.disagg_decode_prealloc_queue.queue) - ) - if self.server_args.disaggregation_decode_enable_offload_kvcache: - queue_size += len(self.decode_offload_manager.ongoing_offload) - - if batch is None and queue_size == 0: + result = self.run_batch(batch) + self.process_batch_result(batch, result) + else: self.self_check_during_idle() self.last_batch = batch @@ -868,7 +826,6 @@ class SchedulerDisaggregationDecodeMixin: def event_loop_overlap_disagg_decode(self: Scheduler): self.result_queue = deque() self.last_batch: Optional[ScheduleBatch] = None - self.last_batch_in_queue = False # last batch is modified in-place, so we need another variable to track if it's extend while True: @@ -878,68 +835,28 @@ class SchedulerDisaggregationDecodeMixin: self.process_decode_queue() batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch - last_batch_in_queue = False batch_result = None if batch: - # Generate fake extend output. - if batch.forward_mode.is_extend(): - # Note: Logprobs should be handled on the prefill engine. - self.stream_output( - batch.reqs, any(req.return_logprob for req in batch.reqs) - ) - trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) - if self.require_mlp_sync: - batch_, batch_result = self._prepare_idle_batch_and_run( - delay_process=True - ) - if batch_: - self.result_queue.append((batch_.copy(), batch_result)) - last_batch_in_queue = True - else: - if self.require_mlp_sync: - self.prepare_mlp_sync_batch(batch) - batch_result = self.run_batch(batch) - self.result_queue.append((batch.copy(), batch_result)) - last_batch_in_queue = True + batch_result = self.run_batch(batch) + self.result_queue.append((batch.copy(), batch_result)) - elif self.require_mlp_sync: - batch, batch_result = self._prepare_idle_batch_and_run( - delay_process=True - ) - if batch: - self.result_queue.append((batch.copy(), batch_result)) - last_batch_in_queue = True - - # Process the results of the previous batch but skip if the last batch is extend - if self.last_batch and self.last_batch_in_queue: + if self.last_batch: tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result(tmp_batch, tmp_result) - - self.launch_batch_sample_if_needed(batch_result) - - queue_size = ( - len(self.waiting_queue) - + len(self.disagg_decode_transfer_queue.queue) - + len(self.disagg_decode_prealloc_queue.queue) - ) - if self.server_args.disaggregation_decode_enable_offload_kvcache: - queue_size += len(self.decode_offload_manager.ongoing_offload) - - if batch is None and queue_size == 0: + elif batch is None: self.self_check_during_idle() + self.launch_batch_sample_if_needed(batch_result) self.last_batch = batch - self.last_batch_in_queue = last_batch_in_queue - def _prepare_idle_batch_and_run(self: Scheduler, delay_process=False): - batch = self.prepare_mlp_sync_batch(None) - result = None - if batch: - result = self.run_batch(batch) - if not delay_process: - self.process_batch_result(batch, result) - return batch, result + def _run_batch_prebuilt_extend( + self: Scheduler, batch: ScheduleBatch + ) -> GenerationBatchResult: + if batch.inner_idle_batch is not None: + return self.run_batch(batch.inner_idle_batch) + + return GenerationBatchResult() def get_next_disagg_decode_batch_to_run( self: Scheduler, @@ -947,7 +864,7 @@ class SchedulerDisaggregationDecodeMixin: """Create fake completed prefill if possible and merge with running batch""" # Merge the prefill batch into the running batch last_batch = self.last_batch - if last_batch and last_batch.forward_mode.is_extend(): + if last_batch and last_batch.forward_mode.is_prebuilt_extend(): # chunked prefill doesn't happen in decode instance. assert self.chunked_req is None # Filter finished batches. @@ -971,6 +888,13 @@ class SchedulerDisaggregationDecodeMixin: self.running_batch = self.update_running_batch(self.running_batch) ret = self.running_batch if not self.running_batch.is_empty() else None + # 1. decode + None -> decode + idle + # 2. decode + prebuilt -> decode + idle (idle forward, prebuilt returns) + # 3. prebuilt + None -> None (None forward, prebuilt returns) + None + # 4. prebuilt + decode + None -> idle (idle forward, prebuilt returns) + decode + idle + if self.require_mlp_sync: + ret = self.prepare_mlp_sync_batch(ret) + if ret: attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()} trace_event_batch("schedule", ret.reqs, attrs=attrs) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 6812397f5..eb40d4be1 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -26,7 +26,7 @@ class ScheduleBatchDisaggregationDecodeMixin: Adapted from .prepare_for_extend(). """ - self.forward_mode = ForwardMode.EXTEND + self.forward_mode = ForwardMode.PREBUILT_EXTEND reqs = self.reqs input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs] extend_num_tokens = sum(len(ids) for ids in input_ids) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 2a26797e9..dc02e3bfb 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1010,6 +1010,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): orig_seq_lens: torch.Tensor = None # shape: [b], int32 # For DP attention + inner_idle_batch: Optional[ScheduleBatch] = None global_num_tokens: Optional[List[int]] = None global_num_tokens_for_logprob: Optional[List[int]] = None is_extend_in_batch: bool = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0738074a7..369aa1410 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1963,6 +1963,10 @@ class Scheduler( for req in batch.reqs: req.time_stats.prefill_start_time = current_time + # Place holder handling for pd-disagg decode event loop + if batch.forward_mode.is_prebuilt_extend(): + return self._run_batch_prebuilt_extend(batch) + # Run forward if self.is_generation: batch_or_worker_batch = batch @@ -2093,10 +2097,10 @@ class Scheduler( if batch.forward_mode.is_decode(): self.process_batch_result_decode(batch, result) trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs) - elif batch.forward_mode.is_extend(): self.process_batch_result_prefill(batch, result) - + elif batch.forward_mode.is_prebuilt_extend(): + self.process_batch_result_prebuilt_extend(batch) elif batch.forward_mode.is_idle(): if self.enable_overlap: if result.copy_done is not None: diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 030e34324..d180de410 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py @@ -30,6 +30,8 @@ class MLPSyncBatchInfo: tp0_info: torch.Tensor = None global_num_tokens: list[int] = None global_num_tokens_for_logprob: list[int] = None + tbo_split_seq_index: torch.Tensor = None + global_forward_mode: int = None def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor: return torch.tensor( @@ -67,6 +69,32 @@ class MLPSyncBatchInfo: self.is_extend_in_batch = bool(tp0_info[:, 3].max().item()) +def _update_gather_batch( + batch: ScheduleBatch, + mlp_sync_info: MLPSyncBatchInfo, + require_mlp_tp_gather: bool, +): + local_batch = batch + + # TODO: handle the case when moe_dense_tp_size != 1 + if not require_mlp_tp_gather: + local_batch.global_num_tokens = [mlp_sync_info.num_tokens] + local_batch.global_num_tokens_for_logprob = [ + mlp_sync_info.num_tokens_for_logprob + ] + else: + local_batch.global_num_tokens = mlp_sync_info.global_num_tokens + local_batch.global_num_tokens_for_logprob = ( + mlp_sync_info.global_num_tokens_for_logprob + ) + local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch + local_batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index + local_batch.global_forward_mode = mlp_sync_info.global_forward_mode + + # Check forward mode for cuda graph + local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph + + def prepare_mlp_sync_batch_raw( local_batch: ScheduleBatch, dp_size: int, @@ -79,7 +107,7 @@ def prepare_mlp_sync_batch_raw( offload_tags: set[str], ): # Check if other DP workers have running batches - if local_batch is None: + if local_batch is None or local_batch.forward_mode.is_prebuilt_extend(): num_tokens = 0 num_tokens_for_logprob = 0 elif local_batch.forward_mode.is_decode(): @@ -101,7 +129,9 @@ def prepare_mlp_sync_batch_raw( num_tokens_for_logprob = local_batch.batch_size() can_cuda_graph = ( - local_batch is None or local_batch.forward_mode.is_decode_or_idle() + local_batch is None + or local_batch.forward_mode.is_decode_or_idle() + or local_batch.forward_mode.is_prebuilt_extend() ) and not disable_cuda_graph is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False @@ -128,29 +158,24 @@ def prepare_mlp_sync_batch_raw( ) mlp_sync_info.all_gather(device=device, group=group) - tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output( - mlp_sync_info.tp0_info[:, 4:6], + mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = ( + tbo_preparer.compute_output( + mlp_sync_info.tp0_info[:, 4:6], + ) ) - if local_batch is None and max(mlp_sync_info.global_num_tokens) > 0: + need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0 + if need_idle_batch and local_batch is None: local_batch = get_idle_batch() - if local_batch is not None: - # TODO: handle the case when moe_dense_tp_size != 1 - if not require_mlp_tp_gather: - local_batch.global_num_tokens = [num_tokens] - local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob] - else: - local_batch.global_num_tokens = mlp_sync_info.global_num_tokens - local_batch.global_num_tokens_for_logprob = ( - mlp_sync_info.global_num_tokens_for_logprob - ) - local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch - local_batch.tbo_split_seq_index = tbo_split_seq_index - local_batch.global_forward_mode = global_forward_mode + batch_to_gather = local_batch + if need_idle_batch and local_batch.forward_mode.is_prebuilt_extend(): + # NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync + local_batch.inner_idle_batch = get_idle_batch() + batch_to_gather = local_batch.inner_idle_batch - # Check forward mode for cuda graph - local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph + if local_batch is not None: + _update_gather_batch(batch_to_gather, mlp_sync_info, require_mlp_tp_gather) return local_batch diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index f665afa90..38f0ce0e8 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import ( RequestStage, ScheduleBatch, ) -from sglang.srt.tracing.trace import trace_slice +from sglang.srt.tracing.trace import trace_slice, trace_slice_batch, trace_slice_end from sglang.srt.utils.common import ceil_div if TYPE_CHECKING: @@ -42,6 +42,26 @@ class SchedulerOutputProcessorMixin: We put them into a separate file to make the `scheduler.py` shorter. """ + def process_batch_result_prebuilt_extend(self: Scheduler, batch: ScheduleBatch): + assert self.disaggregation_mode == DisaggregationMode.DECODE + for req in batch.reqs: + for req in batch.reqs: + req.check_finished() + if req.finished(): + req.time_stats.forward_entry_time = ( + req.time_stats.completion_time + ) = time.perf_counter() + trace_slice_end( + RequestStage.DECODE_QUICK_FINISH, + req.rid, + thread_finish_flag=True, + ) + self.tree_cache.cache_finished_req(req) + + # Note: Logprobs should be handled on the prefill engine. + trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) + self.stream_output(batch.reqs, batch.return_logprob) + def process_batch_result_prefill( self: Scheduler, batch: ScheduleBatch, diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 87418daae..c97090bcb 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -217,6 +217,17 @@ class SchedulerRuntimeCheckerMixin: self.tree_cache.sanity_check() def self_check_during_idle(self: Scheduler): + if self.disaggregation_mode == DisaggregationMode.DECODE: + queue_size = ( + len(self.waiting_queue) + + len(self.disagg_decode_transfer_queue.queue) + + len(self.disagg_decode_prealloc_queue.queue) + ) + if self.server_args.disaggregation_decode_enable_offload_kvcache: + queue_size += len(self.decode_offload_manager.ongoing_offload) + if queue_size: + return + self.check_memory() self.check_tree_cache() self.new_token_ratio = self.init_new_token_ratio diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0c51b2b29..98f129875 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -79,6 +79,10 @@ class ForwardMode(IntEnum): DRAFT_EXTEND_V2 = auto() + # Used in disaggregated decode worker + # Represent a batch of requests having their KV cache ready to start decoding + PREBUILT_EXTEND = auto() + # Split Prefill for PD multiplexing SPLIT_PREFILL = auto() @@ -148,6 +152,9 @@ class ForwardMode(IntEnum): and not self.is_draft_extend() ) + def is_prebuilt_extend(self): + return self == ForwardMode.PREBUILT_EXTEND + @total_ordering class CaptureHiddenMode(IntEnum): diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/two_batch_overlap.py index e72420407..531331c6b 100644 --- a/python/sglang/srt/two_batch_overlap.py +++ b/python/sglang/srt/two_batch_overlap.py @@ -68,7 +68,7 @@ def get_token_num_per_seq( # TODO: may smartly disable TBO when batch size is too small b/c it will slow down def compute_split_seq_index( - forward_mode: "ForwardMode", + forward_mode: ForwardMode, num_tokens: int, extend_lens: Optional[Sequence[int]], token_num_per_seq: Optional[int], @@ -79,7 +79,7 @@ def compute_split_seq_index( elif forward_mode.is_target_verify() or forward_mode.is_decode(): assert token_num_per_seq is not None return (num_tokens // token_num_per_seq) // 2 - elif forward_mode.is_idle(): + elif forward_mode.is_idle() or forward_mode.is_prebuilt_extend(): assert num_tokens == 0 return 0 else: @@ -381,6 +381,8 @@ class TboDPAttentionPreparer: or local_batch.forward_mode.is_decode() ): num_tokens = local_batch.batch_size() * token_num_per_seq + elif local_batch.forward_mode.is_prebuilt_extend(): + num_tokens = 0 else: num_tokens = local_batch.extend_num_tokens self.local_tbo_split_seq_index = compute_split_seq_index( diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 194e4ab17..befce0872 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -158,7 +158,6 @@ suites = { ], "per-commit-4-gpu": [ TestFile("models/test_qwen3_next_models.py", 291), - TestFile("test_disaggregation_dp_attention.py", 155), TestFile("test_gpt_oss_4gpu.py", 300), TestFile("test_local_attn.py", 411), TestFile("test_multi_instance_release_memory_occupation.py", 64), @@ -174,6 +173,7 @@ suites = { TestFile("quant/test_w4a8_deepseek_v3.py", 520), TestFile("test_disaggregation_different_tp.py", 600), TestFile("test_disaggregation_pp.py", 140), + TestFile("test_disaggregation_dp_attention.py", 155), ], "per-commit-4-gpu-b200": [ TestFile("test_deepseek_v3_fp4_4gpu.py", 3600), diff --git a/test/srt/test_disaggregation_dp_attention.py b/test/srt/test_disaggregation_dp_attention.py index c4a90c1c2..f1f02ad60 100644 --- a/test/srt/test_disaggregation_dp_attention.py +++ b/test/srt/test_disaggregation_dp_attention.py @@ -1,4 +1,3 @@ -import os import unittest from types import SimpleNamespace @@ -9,17 +8,21 @@ from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST_MLA, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, popen_launch_pd_server, + try_cached_model, ) class TestDisaggregationDPAttention(TestDisaggregationBase): + PREFILL_DP_SIZE = 4 + DECODE_DP_SIZE = 4 + @classmethod def setUpClass(cls): super().setUpClass() # Temporarily disable JIT DeepGEMM envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False) - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA) # Non blocking start servers cls.start_prefill() @@ -38,9 +41,9 @@ class TestDisaggregationDPAttention(TestDisaggregationBase): "--disaggregation-mode", "prefill", "--tp", - "2", + str(cls.PREFILL_DP_SIZE), "--dp", - "2", + str(cls.PREFILL_DP_SIZE), "--enable-dp-attention", ] prefill_args += cls.transfer_backend + cls.rdma_devices @@ -58,12 +61,13 @@ class TestDisaggregationDPAttention(TestDisaggregationBase): "--disaggregation-mode", "decode", "--tp", - "2", + str(cls.DECODE_DP_SIZE), "--dp", - "2", + str(cls.DECODE_DP_SIZE), "--enable-dp-attention", "--base-gpu-id", - "2", + str(cls.PREFILL_DP_SIZE), + "--prefill-round-robin-balance", ] decode_args += cls.transfer_backend + cls.rdma_devices cls.process_decode = popen_launch_pd_server( @@ -77,7 +81,7 @@ class TestDisaggregationDPAttention(TestDisaggregationBase): args = SimpleNamespace( num_shots=5, data_path=None, - num_questions=200, + num_questions=1400, max_new_tokens=512, parallel=128, host=f"http://{self.base_host}",