[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:
1 parent
68764cc4f4
commit
e37a89ffcb
21 files changed
+804
-29
No files matched your search
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user