[Feature]: Support cfg kv-cache transfer in multi-stage (#1422)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/diffusion/test_diffusion_model_runner.py

@@ -56,7 +56,10 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         enable_cache_dit_summary=enable_cache_dit_summary,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
-    runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_kv_cache=lambda req, target_device=None: None,
+        receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+    )
     return runner
 
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -4,6 +4,7 @@
 stage_args:
   - stage_id: 0
     stage_type: llm
+    prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
       max_batch_size: 1
@@ -39,6 +40,7 @@ stage_args:
       to_stage_1: mooncake_connector
   - stage_id: 1
     stage_type: diffusion
+    cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
     runtime:
       devices: "0"
       max_batch_size: 1

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -4,6 +4,7 @@
 stage_args:
   - stage_id: 0
     stage_type: llm
+    prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
       max_batch_size: 1
@@ -38,6 +39,7 @@ stage_args:
 
   - stage_id: 1
     stage_type: diffusion
+    cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
     runtime:
       devices: "0"
       max_batch_size: 1

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -37,16 +37,16 @@ from vllm_omni.entrypoints.omni import Omni
 # "Generated with seed=52, num_inference_steps=15,
 # prompt='A futuristic city skyline at twilight, cyberpunk style'"
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (68, 107, 134)},
-    {"position": (400, 50), "rgb": (95, 139, 166)},
-    {"position": (700, 100), "rgb": (99, 122, 151)},
-    {"position": (150, 400), "rgb": (111, 125, 153)},
-    {"position": (512, 512), "rgb": (97, 107, 131)},
-    {"position": (700, 400), "rgb": (48, 64, 98)},
-    {"position": (100, 700), "rgb": (79, 63, 84)},
-    {"position": (400, 700), "rgb": (40, 58, 79)},
-    {"position": (700, 700), "rgb": (60, 75, 103)},
-    {"position": (256, 256), "rgb": (97, 128, 156)},
+    {"position": (100, 100), "rgb": (49, 96, 134)},
+    {"position": (400, 50), "rgb": (63, 127, 167)},
+    {"position": (700, 100), "rgb": (70, 101, 141)},
+    {"position": (150, 400), "rgb": (115, 90, 150)},
+    {"position": (512, 512), "rgb": (98, 86, 119)},
+    {"position": (700, 400), "rgb": (29, 42, 91)},
+    {"position": (100, 700), "rgb": (47, 50, 88)},
+    {"position": (400, 700), "rgb": (36, 52, 91)},
+    {"position": (700, 700), "rgb": (45, 58, 99)},
+    {"position": (256, 256), "rgb": (62, 94, 135)},
 ]
 
 # Maximum allowed difference per color channel
@@ -80,6 +80,10 @@ def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_st
     params_list[0].max_tokens = max_tokens  # type: ignore
     if len(params_list) > 1:
         params_list[1].num_inference_steps = num_inference_steps  # type: ignore
+        params_list[1].extra_args = {  # type: ignore
+            "cfg_text_scale": 4.0,
+            "cfg_img_scale": 1.5,
+        }
     return params_list
 
 

tests/entrypoints/test_cfg_companion_tracker.py

@@ -0,0 +1,114 @@
+import time
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def dummy_expand_func(prompt, sp0):
+    if prompt == "expand_me":
+        return [SimpleNamespace(prompt={"prompt": "neg"}, role="cfg_text", request_id_suffix="__cfg_text")]
+    return []
+
+
+@pytest.fixture
+def tracker():
+    sp0 = SimpleNamespace()
+    return CfgCompanionTracker(prompt_expand_func=dummy_expand_func, stage0_sampling_params=sp0, timeout_s=0.1)
+
+
+def test_companion_tracker_initialization(tracker):
+    assert not tracker.is_active
+    assert tracker.num_companions == 0
+
+
+def test_expand_prompts_registers_companions(tracker):
+    request_id_to_prompt = {"req1": "expand_me", "req2": "do_not_expand"}
+
+    pairs = tracker.expand_prompts(request_id_to_prompt)
+
+    assert len(pairs) == 1
+    companion_id, prompt = pairs[0]
+    assert companion_id == "req1__cfg_text"
+    assert prompt == {"prompt": "neg"}
+
+    assert tracker.is_active
+    assert tracker.num_companions == 1
+    assert tracker.is_companion("req1__cfg_text")
+    assert not tracker.is_companion("req2__cfg_text")
+    assert tracker.has_companions("req1")
+    assert not tracker.has_companions("req2")
+
+    comp_map = tracker.get_companion_request_ids("req1")
+    assert comp_map == {"cfg_text": "req1__cfg_text"}
+
+
+def test_companion_lifecycle_success(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    # Defer parent
+    engine_outputs = {"out": 123}
+    tracker.defer_parent("req1", engine_outputs, stage_id=0)
+
+    # Initially not done
+    assert not tracker.all_companions_done("req1")
+
+    # Companion completes
+    parent_id = tracker.on_companion_completed("req1__cfg_text")
+
+    # Parent should be returned since all companions are done and it is pending
+    assert parent_id == "req1"
+    assert tracker.all_companions_done("req1")
+
+    # Pop pending parent
+    popped = tracker.pop_pending_parent("req1")
+    assert popped is not None
+    assert popped["engine_outputs"] == engine_outputs
+    assert popped["stage_id"] == 0
+
+
+def test_companion_lifecycle_failure(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    tracker.defer_parent("req1", {"out": 123}, stage_id=0)
+
+    # Companion fails
+    parent_id, aborted = tracker.on_companion_error("req1__cfg_text")
+
+    assert parent_id == "req1"
+    assert aborted is True
+    assert tracker.is_parent_failed("req1")
+
+    # Parent should be removed from pending list
+    assert tracker.pop_pending_parent("req1") is None
+
+    # Consume failure
+    tracker.consume_parent_failure("req1")
+    assert not tracker.is_parent_failed("req1")
+
+
+def test_companion_lifecycle_timeout(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    tracker.defer_parent("req1", {"out": 123}, stage_id=0)
+
+    # Initially no timeouts
+    timeouts = tracker.check_timeouts()
+    assert len(timeouts) == 0
+
+    # Wait for timeout
+    time.sleep(0.15)
+
+    # Check timeouts again
+    timeouts = tracker.check_timeouts()
+    assert len(timeouts) == 1
+    assert timeouts[0] == "req1"
+
+    # Should be removed from pending
+    assert tracker.pop_pending_parent("req1") is None