[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>
This commit is contained in:
汪志鹏andCursor authored and GitHub committed 2026-03-02 05:06:28 +00:00
1 parent 68764cc4f4
commit e37a89ffcb
21 files changed
+804 -29

No files matched your search

+8
View File
@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
### `prompt_expand_func` (Optional)
A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
### `cfg_kv_collect_func` (Optional)
A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
### `runtime`
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.