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 4, 02:30 UTC
[CI][Elastic EP] Fix Elastic EP Scaling Test Failure (#41792)
Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
haosdent committed
May 9, 2026 at 03:17 UTC
44e6b44a213ca52706c208c639da3fe20d867728
3 files changed
+62
-11
vllm/distributed/elastic_ep/elastic_execute.py
+42
@@ -566,3 +566,45 @@ class ElasticEPScalingExecutor:
566
def prepare_new_worker(self) -> None:
567
with set_current_vllm_config(self.worker.vllm_config):
568
prepare_communication_buffer_for_model(self.worker.model_runner.get_model())
569
+
570
+ def rewarm_workspace(self) -> None:
571
+ # Must run on every DP sibling in lockstep: _dummy_run calls
572
+ # coordinate_batch_across_dp whenever data_parallel_size > 1
573
+ # (gpu_model_runner.py:3663), which deadlocks if any rank skips it.
574
+
575
+ # Save and clear block tables so profile_run/compile_or_warm_up_model
576
+ # don't write dummy slot mappings into real KV-cache blocks (mirrors
577
+ # switch_and_prepare's pattern).
578
+ multi_block_table = self.worker.model_runner.input_batch.block_table
579
+ saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
580
+ for bt in multi_block_table.block_tables:
581
+ saved_block_tables.append(
582
+ (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
583
+ )
584
+ multi_block_table.clear()
585
+
586
+ # _ensure_workspace_size allocates a fresh tensor on grow, leaving
587
+ # captured CUDA graphs with stale data pointers; drop graphs before
588
+ # re-warm so captures realign with the resized buffer.
589
+ self._release_cuda_graphs()
590
+ unlock_workspace()
591
+
592
+ # Grow the MoE workspace at max_num_tokens.
593
+ # compile_or_warm_up_model alone only exercises cudagraph-capture
594
+ # sizes (≤64 tokens for this test) and leaves the workspace at
595
+ # ~10-14 MB; the post-all-to-all per-rank token count under real
596
+ # post-reshuffle routing needs hundreds of MB. Use _dummy_run
597
+ # directly (rather than profile_run) with skip_eplb=True so dummy
598
+ # routing doesn't pollute the just-rebalanced EPLB stats — same
599
+ # convention compile_or_warm_up_model itself uses.
600
+ runner = self.worker.model_runner
601
+ runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
602
+ self.worker.compile_or_warm_up_model()
603
+
604
+ lock_workspace()
605
+
606
+ for bt, (saved_gpu, saved_cpu) in zip(
607
+ multi_block_table.block_tables, saved_block_tables
608
+ ):
609
+ bt.block_table.gpu.copy_(saved_gpu)
610
+ bt.block_table.cpu.copy_(saved_cpu)
vllm/distributed/elastic_ep/elastic_state.py
+5
@@ -538,6 +538,11 @@ class ElasticEPScalingState:
538
self.model_executor.collective_rpc(
539
"elastic_ep_execute", args=("perform_eplb_reshuffle",)
540
)
541
+ # Reshuffle changes per-rank token routing; the locked MoE workspace
542
+ # may now be too small. Rewarm covers both new and existing engines.
543
+ self.model_executor.collective_rpc(
544
+ "elastic_ep_execute", args=("rewarm_workspace",)
545
+ )
546
assert self.new_dp_group is not None
547
if self.new_dp_group.rank() == 0:
548
logger.info("[Elastic EP] EPLB reshuffle completed")
vllm/v1/engine/utils.py
+15
-11
@@ -308,6 +308,18 @@ def get_device_indices(
308
return value
309
310
311
+def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:
312
+ # Ray actor names (RayExecutorV2) and KV-connector engine_ids must
313
+ # be unique across sibling DP engines or registration collides.
314
+ # Use the global DP rank, not a node-local rank, since sibling DP
315
+ # engines can span multiple nodes.
316
+ dp_vllm_config.instance_id = f"{dp_vllm_config.instance_id}_dp{dp_rank}"
317
+ if dp_vllm_config.kv_transfer_config is not None:
318
+ dp_vllm_config.kv_transfer_config.engine_id = (
319
+ f"{dp_vllm_config.kv_transfer_config.engine_id}_dp{dp_rank}"
320
+ )
321
+
322
+
323
class CoreEngineActorManager:
324
"""
325
Utility class to handle creation, readiness, and shutdown
@@ -404,20 +416,10 @@ class CoreEngineActorManager:
416
):
417
dp_vllm_config = copy.deepcopy(vllm_config)
418
if dp_size > 1:
407
- # Append the DP rank to instance_id so that per-engine
408
- # identifiers (e.g. Ray actor names in RayExecutorV2) are
409
- # unique across DP replicas.
410
- dp_vllm_config.instance_id = f"{dp_vllm_config.instance_id}_dp{index}"
419
+ _apply_dp_identity_suffix(dp_vllm_config, index)
420
dp_vllm_config.parallel_config.placement_group = pg
421
local_client = index < local_engine_count
422
414
- if dp_size > 1 and dp_vllm_config.kv_transfer_config is not None:
415
- # modify the engine_id and append the local_dp_rank to it to ensure
416
- # that the kv_transfer_config is unique for each DP rank.
417
- dp_vllm_config.kv_transfer_config.engine_id = (
418
- f"{dp_vllm_config.kv_transfer_config.engine_id}_dp{local_index}"
419
- )
420
-
423
# Ray XPU known issue: dpctl initializes the GPU runtime early, so
424
# setting device env vars in Ray actor's initialization method
425
# will not affect device selection. See:
@@ -789,6 +791,8 @@ class CoreEngineActorManager:
791
for i, (pg, local_rank) in enumerate(zip(placement_groups, local_dp_ranks)):
792
rank = cur_data_parallel_size + i
793
dp_vllm_config = copy.deepcopy(cur_vllm_config)
794
+ if new_data_parallel_size > 1:
795
+ _apply_dp_identity_suffix(dp_vllm_config, rank)
796
dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
797
dp_vllm_config.parallel_config.placement_group = pg
798