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

[Diffusion][HunyuanImage3] Add request paged prefix caching (#7719)

Browse Source 
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
This commit is contained in:
Yukim1 authored and GitHub committed 2026-09-23 15:34:56 +00:00
1 parent 9a1a05facd
commit 205aa01c7b
49 files changed
+2889 -127
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
+51 -9
View File
@@ -13,7 +13,7 @@ from __future__ import annotations
import copy
import gc
import time
from collections.abc import Callable
from collections.abc import Callable, Sequence
from contextlib import AbstractContextManager, nullcontext
from typing import TYPE_CHECKING, Any, cast
@@ -97,9 +97,10 @@ def _dit_any_rank_failed(local_failed: bool) -> bool:
if not torch.distributed.is_initialized():
return local_failed
try:
from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
from vllm_omni.diffusion.distributed import parallel_state
group = get_dit_group()
get_dit_group = getattr(parallel_state, "get_dit_group", None)
group = get_dit_group() if get_dit_group is not None else None
except (AssertionError, ImportError):
group = None
if group is None:
@@ -489,6 +490,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput):
from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load
assert self.od_config.kv_transfer_config is not None
timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0)
if self._kv_receive_progress is not None:
return self._kv_receive_progress.prepare(self._kv_connector, scheduler_output, timeout)
@@ -505,7 +507,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
) -> 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 | tuple[str, int]]) -> int:
def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int:
return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
def refresh_diffusion_kv_block_table_layout(self) -> None:
@@ -538,13 +540,19 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, "
f"active={active_seq_len}, allocated={sequence.seq_len}"
)
# Imported AR KV and local hits have separate owners.
kv_start_pos = (
sequence.num_computed_tokens
if getattr(self.od_config, "kv_transfer_config", None) is not None
else sequence.cached_prefix_len
)
prefill_rows.append(
DiffusionPagedAttentionRow(
request_id=request_metadata.request_id,
sequence_id=sequence.sequence_id,
query_len=sequence.seq_len - sequence.num_computed_tokens,
kv_start_pos=kv_start_pos,
query_len=sequence.seq_len - kv_start_pos,
seq_len=sequence.seq_len,
kv_start_pos=sequence.num_computed_tokens,
)
)
denoise_rows.append(
@@ -800,17 +808,48 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
"Diffusion KV metadata count must match the request batch: "
f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}"
)
for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
native_kv_transfer = getattr(self.od_config, "kv_transfer_config", None) is not None
if native_kv_transfer:
for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
# This field is consumed by Hunyuan only for native
# AR->DiT transfer. Local prefix hits must still run
# their VAE/ViT conditioning path.
req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata)
paged_kv_cached_prefix_len = 0
if not native_kv_transfer:
cached_prefix_lens = {row.kv_start_pos for row in paged_metadata.prefill_rows}
if len(cached_prefix_lens) != 1:
raise ValueError(
"One paged request-level forward requires a uniform cached prefix boundary; "
f"got {sorted(cached_prefix_lens)}"
)
paged_kv_cached_prefix_len = next(iter(cached_prefix_lens))
paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata(
paged_metadata
)
if is_primary:
# Trace the boundary actually passed to the model, not a
# speculative lookup. Useful for warm-cache regressions.
for request_metadata in diffusion_kv_metadata:
for sequence in request_metadata.sequences:
logger.debug(
"Diffusion prefix prefill: request_id=%s sequence_id=%d "
"cached_prefix_len=%d prefix_len=%d query_len=%d",
request_metadata.request_id,
sequence.sequence_id,
sequence.cached_prefix_len,
sequence.prefix_len,
sequence.seq_len - sequence.cached_prefix_len,
)
else:
paged_kv_cached_prefix_len = 0
with (
set_forward_context(
vllm_config=self.vllm_config,
omni_diffusion_config=od_config,
paged_kv_runtime=paged_kv_runtime,
paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
),
paged_kv_context,
@@ -1054,7 +1093,10 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
external_req_id=getattr(sched_new_req.req, "external_req_id", None),
)
if sched_new_req.diffusion_kv_metadata is not None:
if (
sched_new_req.diffusion_kv_metadata is not None
and getattr(self.od_config, "kv_transfer_config", None) is not None
):
new_state.extra["kv_computed_tokens"] = tuple(
seq.num_computed_tokens for seq in sched_new_req.diffusion_kv_metadata.sequences
)