anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Add native paged KV backend for diffusion workers (#6102)

Browse Source 
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Wang fuyinandHongsheng Liu authored and GitHub committed 2026-08-22 21:45:45 +08:00
1 parent c2daa471c3
commit 72ee535fcd
34 files changed
+4595 -71
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/worker/diffusion_model_runner.py
+129 -40
View File
@@ -24,6 +24,7 @@ from vllm.logger import init_logger
from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.cache.cachedit import cache_summary
from vllm_omni.diffusion.cache.prompt_embed_cache import (
install_prompt_embed_cache,
@@ -34,6 +35,11 @@ from vllm_omni.diffusion.compile import regionally_compile
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
DiffusionPagedAttentionRow,
PreparedDiffusionPagedAttentionBatch,
)
from vllm_omni.diffusion.forward_context import set_forward_context
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import (
@@ -144,9 +150,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
self.prompt_embed_cache: Any | None = None
self.input_batch: InputBatch | None = None
self.model_memory_usage = 0
# Rank-local native config generated by the Engine. Physical tensors
# and BlockTable binding remain the responsibility of the paged
# data-plane initialization path.
self.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
vllm_config=vllm_config,
od_config=od_config,
device=device,
)
# Compatibility view for callers that inspect the installed config;
# physical ownership remains in ``diffusion_kv_backend``.
self.kv_cache_config: KVCacheConfig | None = None
# Cache for per-request stepwise state.
@@ -381,41 +391,60 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
"""Collect native specs from cache-enabled loaded attention modules."""
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
return {}
if self.pipeline is None:
raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs")
specs: dict[str, KVCacheSpec] = {}
cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {}
for layer_name, module in self.pipeline.named_modules():
if not isinstance(module, Attention):
continue
spec = module.get_kv_cache_spec(self.vllm_config)
if spec is not None:
specs[layer_name] = spec
if not specs:
cache_layers[layer_name] = (module, spec)
if not cache_layers:
raise RuntimeError(
"paged_scheduler Diffusion KV found no cache-enabled Attention modules "
f"in {type(self.pipeline).__name__}"
)
return specs
return self.diffusion_kv_backend.register_kv_cache_layers(cache_layers)
def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None:
"""Retain the Engine-generated rank-local config for the data plane."""
"""Physically initialize the Engine-generated rank-local config."""
configured_layers = {
layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
}
expected_layers = set(self.get_kv_cache_spec())
if configured_layers != expected_layers:
raise ValueError(
"Rank-local Diffusion KVCacheConfig layer mismatch: "
f"expected={sorted(expected_layers)}, configured={sorted(configured_layers)}"
)
self.kv_cache_config = kv_cache_config
self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config)
self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config
def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata)
def get_diffusion_kv_row(
self,
request_id: str,
sequence_id: int | None,
context_id: str | None = None,
) -> int:
return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id)
def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
def refresh_diffusion_kv_block_table_layout(self) -> None:
self.diffusion_kv_backend.refresh_block_table_layout()
def prepare_paged_attention_batch(
self,
rows: list[DiffusionPagedAttentionRow] | tuple[DiffusionPagedAttentionRow, ...],
) -> PreparedDiffusionPagedAttentionBatch:
"""Prepare model-specific query/write spans against installed Worker rows."""
return self.diffusion_kv_backend.prepare_paged_attention_batch(rows)
def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
"""Expose a prepared paged batch to Omni Attention for one forward."""
return self.diffusion_kv_backend.activate_paged_attention(batch)
def clear_prompt_embed_cache(self) -> None:
"""Evict all cached text-encoder outputs (e.g. between training epochs).
@@ -664,17 +693,25 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
request_id=req.request_id,
metadata=diffusion_kv_metadata,
)
runner_output = self._execute_request_list(
[req],
od_config=self.od_config,
allow_single_output=True,
require_request_batch_support=False,
kv_prefetch_job=kv_prefetch_job,
record_name="pipeline_forward",
)
output = runner_output.runner_outputs[0].result
assert output is not None
return output
installed_request = False
if diffusion_kv_metadata is not None:
self.install_diffusion_kv_metadata(diffusion_kv_metadata)
installed_request = True
try:
runner_output = self._execute_request_list(
[req],
od_config=self.od_config,
allow_single_output=True,
require_request_batch_support=False,
kv_prefetch_job=kv_prefetch_job,
record_name="pipeline_forward",
)
output = runner_output.runner_outputs[0].result
assert output is not None
return output
finally:
if installed_request:
self.remove_diffusion_kv_requests([req.request_id])
def profile_run(self, requests: list[OmniDiffusionRequest]) -> None:
"""Run the maximum per-rank request batch for memory profiling.
@@ -746,14 +783,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
request_id=new_req.req.request_id,
metadata=new_req.diffusion_kv_metadata,
)
reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs]
return self._execute_request_list(
reqs,
od_config=od_config,
allow_single_output=False,
require_request_batch_support=True,
record_name="pipeline_forward_batch",
)
installed_request_ids: list[str] = []
try:
for new_req in scheduler_output.scheduled_new_reqs:
if new_req.diffusion_kv_metadata is not None:
self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata)
installed_request_ids.append(new_req.request_id)
reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs]
return self._execute_request_list(
reqs,
od_config=od_config,
allow_single_output=False,
require_request_batch_support=True,
record_name="pipeline_forward_batch",
)
finally:
if installed_request_ids:
self.remove_diffusion_kv_requests(installed_request_ids)
# ------------------------------------------------------------------
# Step-wise execution
@@ -884,6 +930,40 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
if self.od_config.cache_backend not in (None, "none"):
raise ValueError("Step mode does not support cache_backend yet.")
# Scheduler metadata is installed only for newly admitted requests;
# cached requests continue to use the row installed on their first
# step. Paged rows are retired here for scheduler-finished requests so
# a subsequent wave can reuse the native table slots. Dense execution
# must not issue Worker KV cleanup side effects.
if (
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
is DiffusionKVCacheMode.PAGED_SCHEDULER
and scheduler_output.finished_req_ids
):
self.remove_diffusion_kv_requests(list(scheduler_output.finished_req_ids))
installed_request_ids: list[str] = []
try:
for new_req in scheduler_output.scheduled_new_reqs:
if new_req.diffusion_kv_metadata is not None:
self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata)
installed_request_ids.append(new_req.request_id)
return self._execute_stepwise_core(
scheduler_output,
record_output_peak_memory=record_output_peak_memory,
)
except Exception:
if installed_request_ids:
self.remove_diffusion_kv_requests(installed_request_ids)
raise
def _execute_stepwise_core(
self,
scheduler_output: DiffusionSchedulerOutput,
*,
record_output_peak_memory: bool,
) -> BatchRunnerOutput:
"""Run the denoise step after metadata admission and row installation."""
use_hsdp = self.od_config.parallel_config.use_hsdp
grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
with grad_context:
@@ -1005,7 +1085,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
state.peak_memory_mb,
)
terminal_request_ids = [
runner_output.request_id for runner_output in runner_output_list if runner_output.finished
]
self._update_states_after(states, input_batch, pipeline_interrupted)
if (
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
is DiffusionKVCacheMode.PAGED_SCHEDULER
and terminal_request_ids
):
self.remove_diffusion_kv_requests(terminal_request_ids)
return BatchRunnerOutput.from_list(runner_output_list)