[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Zhou TaichangandClaude Fable 5.1 authored and GitHub committed
1 parent
a9e2117f5e
commit
86fcfe95f1
23 files changed
+1957
-200
No files matched your search
@@ -70,6 +70,7 @@ The lifecycle labels used below are:
|
||||
| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE` | Positive integer; default `32` entries | Each diffusion runner; resolved during model setup | A valid environment value overrides the explicit cache size. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
|
||||
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER` | Boolean with the same accepted spellings as the prompt cache; default disabled | Experimental diffusion session manager; model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
|
||||
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS` | Positive integer; default `64` | Experimental diffusion session manager; model setup | A valid environment value overrides the explicit maximum. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
|
||||
| `VLLM_OMNI_AR_DIFFUSION_KV_GATHER` | `1` enables; default `0` (off) | AR-diffusion KV manager at allocation and attention at dispatch; set before worker startup | Only the exact value `1` enables contiguous KV gathering. History staging additionally requires `ARDiffusionKVConfig.reuse_history_staging`. | Experimental |
|
||||
|
||||
Backend names for `DIFFUSION_ATTENTION_BACKEND` are the members of
|
||||
`DiffusionAttentionBackendEnum`, such as `FLASH_ATTN`, `TORCH_SDPA`,
|
||||
|
||||
@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
"""Make an inventory expansion an explicit review decision."""
|
||||
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
|
||||
assert category_counts == {
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 28,
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 29,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
|
||||
|
||||
Covers the request adapter, the chunk-window spec/manager (registration + the
|
||||
@@ -269,9 +270,14 @@ def test_compiled_paged_write_does_not_clone_the_pool():
|
||||
# drives the reinplace pass, so a drift here would silently test nothing.
|
||||
real_schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
|
||||
real_mutated = [arg.name for arg in real_schema.arguments if arg.alias_info and arg.alias_info.is_write]
|
||||
assert real_mutated == ["key_pool", "value_pool"], (
|
||||
# The staging buffers are mutated too when reuse_history_staging is on, but they are separate
|
||||
# allocations and not what this test is about; the pools are what must not be cloned.
|
||||
assert real_mutated[:2] == ["key_pool", "value_pool"], (
|
||||
f"the real paged-write op now mutates {real_mutated}; this test still models two pools"
|
||||
)
|
||||
assert set(real_mutated) - {"key_pool", "value_pool"} <= {"stage_key", "stage_value"}, (
|
||||
f"unexpected mutated argument in the real paged-write op: {real_mutated}"
|
||||
)
|
||||
|
||||
num_blocks, heads, dim = 8, 4, 16
|
||||
cache_shape = (num_blocks, BLOCK, heads, dim)
|
||||
@@ -357,6 +363,7 @@ def _make_tiny_capacity_kv(
|
||||
available_bytes: int,
|
||||
gpu_memory_fraction: float = 1.0,
|
||||
model_owned_state_bytes_per_session: int = 0,
|
||||
reuse_history_staging: bool = False,
|
||||
) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
@@ -365,6 +372,7 @@ def _make_tiny_capacity_kv(
|
||||
window_chunks=3,
|
||||
sink_chunks=3,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
reuse_history_staging=reuse_history_staging,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
@@ -382,6 +390,33 @@ def _make_tiny_capacity_kv(
|
||||
)
|
||||
|
||||
|
||||
def test_history_staging_is_reserved_in_the_kv_budget(monkeypatch):
|
||||
"""The per-layer staging pair is worker-wide memory allocated after admission; the budget must hold it."""
|
||||
monkeypatch.setenv("VLLM_OMNI_AR_DIFFUSION_KV_GATHER", "1")
|
||||
# Same geometry as the capacity-two test: 192 bytes fits two sessions exactly without staging.
|
||||
plain = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
|
||||
assert plain.session_capacity == 2 and plain.history_staging_reserved_bytes == 0
|
||||
|
||||
staged = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192, reuse_history_staging=True)
|
||||
# 2 (K, V) x 1 layer x (sink 3 + window 3 + one action-capacity block) tokens x 1 head x 1 dim x 4 bytes.
|
||||
assert staged.history_staging_reserved_bytes == 2 * 1 * 7 * 1 * 1 * 4
|
||||
# It comes out of the same budget, so the second session no longer fits.
|
||||
assert staged.session_capacity == 1
|
||||
assert staged.num_blocks_total * 8 + staged.history_staging_reserved_bytes <= 192
|
||||
|
||||
# Boundary: one session fits at 128 bytes without staging; with staging the same budget is rejected
|
||||
# up front instead of failing on the first forward.
|
||||
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=128)
|
||||
with pytest.raises(ValueError, match="cannot fit one session"):
|
||||
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=128, reuse_history_staging=True)
|
||||
|
||||
|
||||
def test_history_staging_is_not_reserved_when_the_gather_path_is_off(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_OMNI_AR_DIFFUSION_KV_GATHER", raising=False)
|
||||
staged = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192, reuse_history_staging=True)
|
||||
assert staged.history_staging_reserved_bytes == 0 and staged.session_capacity == 2
|
||||
|
||||
|
||||
def test_capacity_two_retains_both_windows_and_allocates_next_block():
|
||||
# LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
|
||||
# 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
|
||||
@@ -588,3 +623,70 @@ def test_non_contiguous_branch_indices_rejected():
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 1),),
|
||||
session_capacity=1,
|
||||
)
|
||||
|
||||
|
||||
def test_staged_window_reuse_matches_a_full_gather():
|
||||
"""Reusing the staged history must be byte-identical to re-gathering the whole window.
|
||||
|
||||
This is the property the optimisation rests on: the leading blocks were staged by an earlier forward of
|
||||
the same AR block, so skipping them can only be correct if what is already there equals what a fresh
|
||||
gather would write. The test stages once, mutates only the current chunk's blocks in the pool (what a
|
||||
later probe of the same block does), restages the tail alone, and compares against a full gather.
|
||||
"""
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import _stage_window
|
||||
|
||||
torch.manual_seed(0)
|
||||
num_blocks, block_size, heads, dim = 6, 4, 2, 8
|
||||
current_blocks = 2
|
||||
key_cache = torch.randn(num_blocks, block_size, heads, dim)
|
||||
value_cache = torch.randn(num_blocks, block_size, heads, dim)
|
||||
block_ids = torch.tensor([5, 4, 3, 2, 1, 0])
|
||||
stage_key = torch.zeros(num_blocks * block_size, heads, dim)
|
||||
stage_value = torch.zeros(num_blocks * block_size, heads, dim)
|
||||
|
||||
_stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=0)
|
||||
|
||||
# A later probe rewrites only the current chunk's blocks in the pool.
|
||||
for slot in block_ids[num_blocks - current_blocks :]:
|
||||
key_cache[slot] = torch.randn(block_size, heads, dim)
|
||||
value_cache[slot] = torch.randn(block_size, heads, dim)
|
||||
|
||||
_stage_window(
|
||||
stage_key,
|
||||
stage_value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
block_ids,
|
||||
num_blocks,
|
||||
block_size,
|
||||
first_block=num_blocks - current_blocks,
|
||||
)
|
||||
|
||||
full_key = torch.zeros_like(stage_key)
|
||||
full_value = torch.zeros_like(stage_value)
|
||||
_stage_window(full_key, full_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=0)
|
||||
|
||||
torch.testing.assert_close(stage_key, full_key, rtol=0, atol=0)
|
||||
torch.testing.assert_close(stage_value, full_value, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_staged_window_tail_refresh_leaves_history_untouched():
|
||||
"""The tail restage must not write the history rows -- that is what makes it cheaper."""
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import _stage_window
|
||||
|
||||
torch.manual_seed(1)
|
||||
num_blocks, block_size, heads, dim = 5, 4, 1, 4
|
||||
key_cache = torch.randn(num_blocks, block_size, heads, dim)
|
||||
value_cache = torch.randn(num_blocks, block_size, heads, dim)
|
||||
block_ids = torch.arange(num_blocks)
|
||||
stage_key = torch.full((num_blocks * block_size, heads, dim), -1.0)
|
||||
stage_value = torch.full((num_blocks * block_size, heads, dim), -1.0)
|
||||
|
||||
_stage_window(
|
||||
stage_key, stage_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=num_blocks - 1
|
||||
)
|
||||
|
||||
history_rows = (num_blocks - 1) * block_size
|
||||
assert (stage_key[:history_rows] == -1.0).all(), "history rows were rewritten by a tail restage"
|
||||
assert (stage_value[:history_rows] == -1.0).all()
|
||||
torch.testing.assert_close(stage_key[history_rows:], key_cache[-1].reshape(block_size, heads, dim), rtol=0, atol=0)
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Tests for AR-Diffusion paged self-attention contexts."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -20,6 +21,8 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
|
||||
ar_diffusion_paged_attention,
|
||||
paged_write_attn,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -32,8 +35,12 @@ POS = "positive"
|
||||
NEG = "negative"
|
||||
|
||||
|
||||
def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
|
||||
cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
|
||||
def make_state(
|
||||
*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
|
||||
):
|
||||
cfg = ARDiffusionKVConfig(
|
||||
enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
|
||||
)
|
||||
kv = ARDiffusionKVCache(
|
||||
cfg,
|
||||
num_layers=num_layers,
|
||||
@@ -236,6 +243,86 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
|
||||
assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
|
||||
def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks):
|
||||
"""A second probe of the same AR block restages its current K/V where the table actually holds it.
|
||||
|
||||
The padded table always ends in at least one action-capacity block, and while the window is still
|
||||
growing in unused window capacity too, so "the last current_blocks entries" is padding. The refresh
|
||||
has to start right after the visible history: empty (window growing), partial, full, and after a slide.
|
||||
"""
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
device = torch.device("cpu")
|
||||
dtype = torch.float32
|
||||
kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True)
|
||||
if history_chunks:
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
|
||||
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.ensure_video_slots(device)
|
||||
key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
|
||||
|
||||
def probe(step: int):
|
||||
# Each probe of the block writes new current K/V, then stages.
|
||||
kv._k_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), float(step))
|
||||
kv._v_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), -float(step))
|
||||
# Store the metadata on the context the way prepare() does.
|
||||
(ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = (
|
||||
ctx.build_block_table(action_len=0, query_len=BLOCK, device=device)
|
||||
)
|
||||
block_table, max_seq_len = ctx.block_table, ctx.max_seq_len
|
||||
ctx._prepare_history_staging(0)
|
||||
n_blocks = max_seq_len // BLOCK
|
||||
block_ids = block_table[0, :n_blocks].to(torch.long)
|
||||
stage_k, stage_v = ctx.history_staging(0)
|
||||
paged_attention_module._stage_window(
|
||||
stage_k,
|
||||
stage_v,
|
||||
key_cache,
|
||||
value_cache,
|
||||
block_ids,
|
||||
n_blocks,
|
||||
BLOCK,
|
||||
first_block=ctx.stage_first_block if ctx.reuse_history else 0,
|
||||
)
|
||||
full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
||||
full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
||||
return stage_k, stage_v, full_k, full_v, n_blocks
|
||||
|
||||
stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
|
||||
assert ctx.reuse_history is False
|
||||
assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
|
||||
|
||||
stage_k, stage_v, full_k, full_v, n_blocks = probe(2)
|
||||
assert ctx.reuse_history is True
|
||||
visible_history_blocks = min(history_chunks, 2 - 1) # window of 2 blocks minus the current one
|
||||
assert ctx.stage_first_block == visible_history_blocks
|
||||
# The padded table is wider than the live window, so the end-of-table guess is not the offset.
|
||||
assert n_blocks - 1 != ctx.stage_first_block
|
||||
assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gather_enabled", [False, True])
|
||||
def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
|
||||
if gather_enabled:
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
else:
|
||||
monkeypatch.delenv(KV_GATHER_ENV, raising=False)
|
||||
device = torch.device("cpu")
|
||||
kv, st = make_state(device=device, reuse_history_staging=True)
|
||||
# The manager owns the pairs and allocates them only when their consumer is on.
|
||||
assert bool(kv.history_staging) is gather_enabled
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.ensure_video_slots(device)
|
||||
(ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = ctx.build_block_table(
|
||||
action_len=0, query_len=BLOCK, device=device
|
||||
)
|
||||
ctx._prepare_history_staging(0)
|
||||
assert ctx.staging_enabled is gather_enabled
|
||||
stage_k, _ = ctx.history_staging(0)
|
||||
assert (stage_k is not None) is gather_enabled
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
||||
@pytest.mark.parametrize("history_chunks", [1, 3])
|
||||
@pytest.mark.parametrize("action_len", [0, 3])
|
||||
@@ -376,13 +463,27 @@ def test_custom_op_registration_idempotent():
|
||||
assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
|
||||
|
||||
|
||||
def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
|
||||
def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
|
||||
# Older PyTorch ADInplaceOrView handlers index positional mutable inputs
|
||||
# directly. Default-valued trailing inputs can disappear before that handler.
|
||||
schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
|
||||
mutable = [arg for arg in schema.arguments if arg.alias_info is not None and arg.alias_info.is_write]
|
||||
assert {arg.name for arg in mutable} == {"key_pool", "value_pool", "stage_key", "stage_value"}
|
||||
assert all(not arg.has_default_value() for arg in mutable)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reuse_history_staging", [False, True])
|
||||
def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
|
||||
"""The op must trace as one opaque node: fullgraph OK, and changed tensor
|
||||
VALUES (new slots / block ids) must not trigger recompilation."""
|
||||
import torch._dynamo
|
||||
|
||||
device = torch.device("cpu")
|
||||
kv, st = make_state(num_layers=2, window_chunks=2)
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
kv, st = make_state(num_layers=2, window_chunks=2, reuse_history_staging=reuse_history_staging)
|
||||
if reuse_history_staging:
|
||||
# Hold the host-side staging offset fixed while table/slot values change.
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=2, dtype=torch.float32, device=device)
|
||||
|
||||
def run_one_forward(commit):
|
||||
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit)
|
||||
@@ -391,9 +492,8 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
|
||||
q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
||||
k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
||||
v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
||||
# Both layers through ONE compiled fn: layer_idx is a tensor, so a
|
||||
# different layer must NOT recompile (all 40 DiT blocks share the
|
||||
# block-forward code object in production).
|
||||
# Both layers use one compiled function; the tensor-valued layer index
|
||||
# must not specialize it.
|
||||
for layer_ctx in contexts:
|
||||
out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
|
||||
st.commit_paged_context(POS)
|
||||
@@ -419,3 +519,52 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
|
||||
# Leave a clean dynamo state for later suites in the same pytest process
|
||||
# (e.g. model_executor transformers models).
|
||||
torch._dynamo.reset()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
||||
@pytest.mark.parametrize("history_chunks", [0, 1, 3])
|
||||
@pytest.mark.parametrize("action_len", [0, 3])
|
||||
def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
|
||||
"""VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1 gathers the visible blocks and runs varlen FA3 without a block table.
|
||||
|
||||
Covers an empty, partial and full window plus a partially filled action block:
|
||||
the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
|
||||
"""
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.float16
|
||||
kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
|
||||
if history_chunks:
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
|
||||
|
||||
layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
|
||||
ctx = layer_ctx.forward_ctx
|
||||
current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
||||
current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
||||
action_k = action_v = None
|
||||
if action_len:
|
||||
action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
||||
action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
||||
query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
||||
ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
|
||||
inputs = layer_ctx.to_layer_inputs()
|
||||
|
||||
def run() -> torch.Tensor:
|
||||
return paged_write_attn(
|
||||
inputs,
|
||||
query[0],
|
||||
current_k[0],
|
||||
current_v[0],
|
||||
action_k[0] if action_len else None,
|
||||
action_v[0] if action_len else None,
|
||||
HEAD_DIM**-0.5,
|
||||
).unsqueeze(0)
|
||||
|
||||
monkeypatch.delenv(KV_GATHER_ENV, raising=False)
|
||||
paged = run()
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
gathered = run()
|
||||
assert torch.isfinite(gathered).all()
|
||||
# Same kernel family on the same K/V: only accumulation order differs.
|
||||
torch.testing.assert_close(gathered, paged, rtol=2e-3, atol=2e-3)
|
||||
@@ -461,3 +461,46 @@ def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
|
||||
device_kind="cuda",
|
||||
master_port=29501,
|
||||
)
|
||||
|
||||
|
||||
def _run_fused_qkv_matches_three_4d(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
device_kind: DeviceKind,
|
||||
master_port: int,
|
||||
) -> None:
|
||||
"""LingBot's fused (B, S/N, 3, H, D) exchange must equal three per-tensor 4D exchanges."""
|
||||
from vllm_omni.diffusion.distributed.comm import all_to_all_5D
|
||||
|
||||
device = _worker_device(local_rank, device_kind)
|
||||
if device_kind == "cuda":
|
||||
current_omni_platform.set_device(device)
|
||||
_init_worker(local_rank, world_size, master_port, device_kind)
|
||||
initialize_model_parallel(ulysses_degree=world_size)
|
||||
sp_group = get_sp_group().ulysses_group
|
||||
try:
|
||||
batch, seq_per_rank, heads, head_size = 1, 6, 4 * world_size, 8
|
||||
# Distinguishable values: rank, q/k/v index, position, head, feature.
|
||||
base = torch.arange(seq_per_rank * heads * head_size, dtype=torch.float32, device=device).reshape(
|
||||
batch, seq_per_rank, heads, head_size
|
||||
)
|
||||
tensors = [base + 1000.0 * local_rank + 100000.0 * which for which in range(3)]
|
||||
separate = [SeqAllToAll4D.apply(sp_group, t, 2, 1, False) for t in tensors]
|
||||
fused = all_to_all_5D(torch.stack(tensors, dim=2), scatter_idx=3, gather_idx=1, group=sp_group)
|
||||
assert fused.shape == (batch, seq_per_rank * world_size, 3, heads // world_size, head_size)
|
||||
for which, expected in enumerate(separate):
|
||||
assert torch.equal(fused.unbind(2)[which], expected), f"q/k/v index {which} differs on rank {local_rank}"
|
||||
finally:
|
||||
destroy_distributed_env()
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
|
||||
torch.multiprocessing.spawn(
|
||||
_run_fused_qkv_matches_three_4d,
|
||||
args=(world_size, "cpu", 29613),
|
||||
nprocs=world_size,
|
||||
)
|
||||
@@ -356,6 +356,94 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
|
||||
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
|
||||
|
||||
def _reference_conv_input(conv, x, cache_x):
|
||||
"""The conv input the previous implementation built: cat, F.pad, then halo_exchange."""
|
||||
padding = list(conv._padding)
|
||||
if cache_x is not None and padding[4] > 0:
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
x = torch.nn.functional.pad(x, padding)
|
||||
x_padded, _, _ = wan_spatial_shard.halo_exchange(
|
||||
x, group=object(), halo_size=conv.halo_size, split_dim=conv.split_dim
|
||||
)
|
||||
return x_padded, x.shape[conv.split_tensor_dim]
|
||||
|
||||
|
||||
def _dist_conv(split_dim, kernel=(3, 3, 3), padding=(1, 1, 1)):
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
|
||||
|
||||
torch.manual_seed(0)
|
||||
source = WanCausalConv3d(3, 4, kernel, padding=padding)
|
||||
return source, wan_spatial_shard.WanDistCausalConv3d(source, object(), split_dim=split_dim)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("split_dim", ["width", "height"])
|
||||
@pytest.mark.parametrize("with_cache", [False, True])
|
||||
def test_single_pass_conv_input_matches_cat_pad_halo_on_one_rank(monkeypatch, split_dim, with_cache):
|
||||
"""One rank, halo-bearing kernel: same bytes as the previous cat -> F.pad -> halo_exchange path."""
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
|
||||
_, conv = _dist_conv(split_dim)
|
||||
x = torch.randn(1, 3, 2, 6, 8)
|
||||
cache_x = torch.randn(1, 3, 2, 6, 8) if with_cache else None
|
||||
|
||||
expected, expected_extent = _reference_conv_input(conv, x, cache_x)
|
||||
got, got_extent = conv._assemble_input(x, cache_x)
|
||||
|
||||
assert got.shape == expected.shape and got_extent == expected_extent
|
||||
assert torch.equal(got, expected)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_single_pass_conv_rejects_cache_longer_than_causal_padding():
|
||||
_, conv = _dist_conv("width")
|
||||
with pytest.raises(AssertionError, match="temporal cache exceeds"):
|
||||
conv._assemble_input(torch.randn(1, 3, 2, 6, 8), torch.randn(1, 3, 3, 6, 8))
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("split_dim", ["width", "height"])
|
||||
def test_single_pass_conv_without_a_halo_matches_the_stock_causal_conv(monkeypatch, split_dim):
|
||||
"""A kernel with no spatial extent along the split (the temporal conv) needs no halo and equals the stock conv."""
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
|
||||
source, conv = _dist_conv(split_dim, kernel=(3, 1, 1), padding=(1, 0, 0))
|
||||
x = torch.randn(1, 3, 2, 6, 8)
|
||||
cache_x = torch.randn(1, 3, 2, 6, 8)
|
||||
|
||||
assert torch.equal(conv(x, cache_x), source(x, cache_x))
|
||||
assert torch.equal(conv(x, None), source(x, None))
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("split_dim", ["width", "height"])
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
@pytest.mark.parametrize("kernel,padding", [((3, 3, 3), (1, 1, 1)), ((3, 1, 1), (0, 0, 0))])
|
||||
def test_single_pass_conv_input_matches_cat_pad_halo_on_two_ranks(monkeypatch, split_dim, rank, kernel, padding):
|
||||
"""Same bytes as cat -> F.pad -> halo_exchange, halos included, on either rank of a pair."""
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
|
||||
|
||||
def fake_p2p(*, rank, world_size, group, top_row_ref, bottom_row_ref, recv_top_buf, recv_bottom_buf):
|
||||
# A neighbour's halo is whatever it sent; stand in with a marker derived from our own edge so
|
||||
# both the reference and the single-pass path see the same exchange. Global edges stay zero.
|
||||
recv_top_buf.copy_(top_row_ref * 0 + 7.0) if rank > 0 else recv_top_buf.zero_()
|
||||
recv_bottom_buf.copy_(bottom_row_ref * 0 + 9.0) if rank < world_size - 1 else recv_bottom_buf.zero_()
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_halo_exchange_p2p", fake_p2p)
|
||||
_, conv = _dist_conv(split_dim, kernel=kernel, padding=padding)
|
||||
x = torch.randn(1, 3, 1, 6, 8)
|
||||
cache_x = torch.randn(1, 3, 2, 6, 8)
|
||||
|
||||
expected, expected_extent = _reference_conv_input(conv, x, cache_x)
|
||||
got, got_extent = conv._assemble_input(x, cache_x)
|
||||
|
||||
assert got.shape == expected.shape and got_extent == expected_extent
|
||||
assert torch.equal(got, expected)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@@ -74,6 +74,17 @@ def _install_vllm_stubs() -> None:
|
||||
_SeqAllToAll4D,
|
||||
)
|
||||
|
||||
def _all_to_all_5D(value, scatter_idx=3, gather_idx=1, group=None, use_sync=False):
|
||||
# Single-rank stub: the fused q/k/v exchange is the identity when N == 1.
|
||||
del scatter_idx, gather_idx, group, use_sync
|
||||
return value
|
||||
|
||||
setattr(
|
||||
sys.modules["vllm_omni.diffusion.distributed.comm"],
|
||||
"all_to_all_5D",
|
||||
_all_to_all_5D,
|
||||
)
|
||||
|
||||
def get_sp_group():
|
||||
return SimpleNamespace(
|
||||
ulysses_world_size=1,
|
||||
@@ -588,26 +599,7 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP
|
||||
torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
|
||||
|
||||
|
||||
def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch):
|
||||
module = _load_module()
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_sp_group",
|
||||
lambda: SimpleNamespace(
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=1,
|
||||
ulysses_group=None,
|
||||
),
|
||||
)
|
||||
attention = module.LingBotCrossAttention(dim=8, num_heads=4)
|
||||
full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2)
|
||||
shard = attention.shard_kv_heads(full)
|
||||
torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0)
|
||||
assert shard.is_contiguous()
|
||||
assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"])
|
||||
@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention"])
|
||||
def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
|
||||
module = _load_module()
|
||||
monkeypatch.setattr(
|
||||
@@ -621,3 +613,35 @@ def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
|
||||
)
|
||||
with pytest.raises(ValueError, match="heads must be divisible"):
|
||||
getattr(module, attention_class)(dim=8, num_heads=4)
|
||||
|
||||
|
||||
def test_supplied_camera_modulation_matches_computing_it_inline() -> None:
|
||||
"""Passing a prebuilt modulation must be arithmetically identical to letting the block build its own."""
|
||||
module = _load_module()
|
||||
torch.manual_seed(0)
|
||||
block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
|
||||
camera = torch.randn(1, 3, 4)
|
||||
|
||||
scale, shift = block.camera_modulation(camera)
|
||||
inline_scale, inline_shift = block.camera_modulation(camera)
|
||||
|
||||
torch.testing.assert_close(scale, inline_scale, rtol=0, atol=0)
|
||||
torch.testing.assert_close(shift, inline_shift, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_the_camera_modulation_reads_only_the_camera_tokens() -> None:
|
||||
"""The reuse window is only sound because the modulation depends on nothing else that varies per forward."""
|
||||
module = _load_module()
|
||||
torch.manual_seed(0)
|
||||
block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
|
||||
camera = torch.randn(1, 3, 4)
|
||||
|
||||
first = block.camera_modulation(camera)
|
||||
# Same camera, different everything else a forward would carry: the modulation cannot move.
|
||||
second = block.camera_modulation(camera.clone())
|
||||
|
||||
torch.testing.assert_close(first[0], second[0], rtol=0, atol=0)
|
||||
torch.testing.assert_close(first[1], second[1], rtol=0, atol=0)
|
||||
# A different camera must move it, or the cache key would be hiding a real dependency.
|
||||
other = block.camera_modulation(camera + 1)
|
||||
assert not torch.equal(first[0], other[0])
|
||||
@@ -837,6 +837,9 @@ def _rollout(model, mode, dtype, batch):
|
||||
frames_per_block=_FRAMES,
|
||||
max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME,
|
||||
cross_attention_lengths={"text": 5},
|
||||
# Cross-attention keeps every local head on every SP rank, so its pool is wider than the
|
||||
# self-attention share above.
|
||||
cross_attention_kv_heads={"text": model.blocks[0].cross_attn.num_local_heads},
|
||||
device=device,
|
||||
)
|
||||
state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS)
|
||||
@@ -896,7 +899,6 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_distributed_env,
|
||||
destroy_model_parallel,
|
||||
get_sp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
@@ -948,11 +950,10 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
|
||||
bound = 1e-5 if dtype == torch.float32 else 1e-2
|
||||
error = (result - expected).double().norm() / expected.double().norm()
|
||||
assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}"
|
||||
local_heads = _HEADS // tp // sp
|
||||
first = (
|
||||
get_tensor_model_parallel_rank() * (_HEADS // tp)
|
||||
+ get_sp_group().ulysses_rank * local_heads
|
||||
)
|
||||
# Cross-attention K/V is replicated across the Ulysses group, not split over it, so
|
||||
# every SP rank holds its whole TP shard of heads. Only TP slices the cache now.
|
||||
local_heads = _HEADS // tp
|
||||
first = get_tensor_model_parallel_rank() * local_heads
|
||||
for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True):
|
||||
tolerance = 1e-5 if dtype == torch.float32 else 2e-2
|
||||
torch.testing.assert_close(
|
||||
@@ -1033,3 +1034,56 @@ def test_temporal_rope_extension_keeps_spatial_limit():
|
||||
model._rotary_embedding(
|
||||
frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_the_camera_injector_runs_once_per_held_cache_and_gives_the_same_output() -> None:
|
||||
"""Two forwards sharing one caller-held cache: the injector runs once per block, outputs unchanged."""
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=2).eval()
|
||||
hidden_states = torch.randn(1, 36, 1, 4, 4)
|
||||
timestep = torch.tensor([1.0])
|
||||
encoder_hidden_states = torch.randn(1, 3, 6)
|
||||
camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
|
||||
|
||||
def run(cache, camera_cache=None):
|
||||
return model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
camera_modulation_cache=camera_cache,
|
||||
)
|
||||
|
||||
# Baseline: no cache, so every forward rebuilds the injector in every block.
|
||||
uncached_first = run(_cache(module, model))
|
||||
uncached_second = run(_cache(module, model))
|
||||
injector_calls = [attention_tests._record_outputs(block.cam_injector_layer1) for block in model.blocks]
|
||||
held = module.CameraModulationCache()
|
||||
cached_first = run(_cache(module, model), held)
|
||||
cached_second = run(_cache(module, model), held)
|
||||
|
||||
# One build per block for both forwards, instead of one per block per forward.
|
||||
assert [len(calls) for calls in injector_calls] == [1, 1]
|
||||
torch.testing.assert_close(cached_first, uncached_first, rtol=0, atol=0)
|
||||
torch.testing.assert_close(cached_second, uncached_second, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_timestep_projection_is_staged_into_one_static_buffer_per_shape() -> None:
|
||||
"""Two forwards' projections land in the same buffer; the blocks never see the fresh per-forward tensor."""
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1).eval()
|
||||
first = torch.randn(1, 8, 6, model.dim)
|
||||
second = torch.randn(1, 8, 6, model.dim)
|
||||
|
||||
staged_first = model._stage_timestep_projection(first)
|
||||
assert staged_first is not first and torch.equal(staged_first, first)
|
||||
staged_second = model._stage_timestep_projection(second)
|
||||
assert staged_second is staged_first and torch.equal(staged_second, second)
|
||||
# A different shape gets its own buffer; the first one is kept.
|
||||
other = model._stage_timestep_projection(torch.randn(1, 4, 6, model.dim))
|
||||
assert other is not staged_first and len(model._timestep_projection_buffers) == 2
|
||||
@@ -4,11 +4,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import os
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
|
||||
@@ -29,6 +31,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import (
|
||||
)
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
from vllm_omni.experimental.ar_diffusion.capability import supports_chunk_step_grouping
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
@@ -127,6 +130,10 @@ class _RecordingTransformer(nn.Module):
|
||||
sink_size=3,
|
||||
)
|
||||
self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()])
|
||||
for block in self.blocks:
|
||||
# The spec reads head geometry from both attentions; cross-attention keeps every local head.
|
||||
block.self_attn = SimpleNamespace(num_sp_heads=2)
|
||||
block.cross_attn = SimpleNamespace(num_local_heads=2)
|
||||
self.calls: list[dict] = []
|
||||
self.cache_allocations: list[dict] = []
|
||||
self.raise_on_call = raise_on_call
|
||||
@@ -146,6 +153,7 @@ class _RecordingTransformer(nn.Module):
|
||||
"cache_id": id(kwargs["cache"]),
|
||||
"start_frame": kwargs["start_frame"],
|
||||
"update_cache": kwargs["update_cache"],
|
||||
"camera_cache": kwargs.get("camera_modulation_cache"),
|
||||
}
|
||||
self.calls.append(call)
|
||||
if self.raise_on_call == len(self.calls):
|
||||
@@ -534,14 +542,15 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None:
|
||||
assert spec.sink_frames == 3
|
||||
assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)]
|
||||
assert spec.cross_attention_lengths == {"text": 512}
|
||||
# One condition block, committed + pending encoder caches, and decoder state.
|
||||
# One condition block, committed + pending encoder caches, the session's
|
||||
# cached prompt embedding (512 tokens x text_dim 8 x fp32), and decoder state.
|
||||
# The stub encoder retains 2 RGB frames at 16x16 and 2 feature frames at 2x2.
|
||||
encoder_bytes = (2 * 3 * 16 * 16 + 2 * 16 * 2 * 2) * 4
|
||||
assert spec.model_owned_state_bytes_per_session == 960 + 2 * encoder_bytes + 2_421_248
|
||||
assert spec.model_owned_state_bytes_per_session == 960 + 384 + 2 * encoder_bytes + 2_421_248
|
||||
|
||||
|
||||
def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
|
||||
"""Admission includes encoder state even when the decoder cannot stream."""
|
||||
"""Admission includes encoder state and the prompt embedding even when the decoder cannot stream."""
|
||||
module = _load_pipeline_module()
|
||||
|
||||
def spec_pipeline():
|
||||
@@ -560,13 +569,29 @@ def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
|
||||
# than the image condition alone could account for.
|
||||
assert widened - small > 2_000_000
|
||||
|
||||
# Removing decoder support leaves the condition and both encoder histories.
|
||||
# Removing decoder support leaves the condition, both encoder histories and
|
||||
# the prompt-embedding cache.
|
||||
bare = spec_pipeline()
|
||||
for attribute in ("decoder", "post_quant_conv"):
|
||||
if hasattr(bare.vae, attribute):
|
||||
delattr(bare.vae, attribute)
|
||||
bare.vae._cached_conv_counts.pop("decoder")
|
||||
assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 2 * 6_656
|
||||
assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 384 + 2 * 6_656
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sp_size", [1, 4])
|
||||
def test_session_admission_budgets_camera_modulation_per_rank(sp_size: int) -> None:
|
||||
pipeline = _pipeline(_load_pipeline_module())
|
||||
pipeline._ar_height, pipeline._ar_width = 480, 832
|
||||
pipeline.od_config.parallel_config.ulysses_degree = sp_size
|
||||
pipeline.transformer.config.num_layers = 40
|
||||
pipeline.transformer.config.num_attention_heads = 40
|
||||
pipeline.transformer.config.attention_head_dim = 128
|
||||
pipeline.transformer._dtype = torch.bfloat16
|
||||
forty_layers = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
|
||||
pipeline.transformer.config.num_layers = 1
|
||||
one_layer = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
|
||||
assert forty_layers - one_layer == 2 * 39 * (4680 // sp_size) * 5120 * 2
|
||||
|
||||
|
||||
def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
|
||||
@@ -1672,7 +1697,6 @@ def test_typed_ticks_generate_one_global_block_and_return_standard_metadata(
|
||||
"applied_event_ids": [1],
|
||||
}
|
||||
assert pipeline._ar_sessions["world-1"].next_chunk_index == 2
|
||||
assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley"
|
||||
assert pipeline._ar_sessions["world-1"].camera_tail is not None
|
||||
assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4)
|
||||
expected_generator = torch.Generator(device="cpu").manual_seed(17)
|
||||
@@ -1846,8 +1870,11 @@ def test_typed_ticks_keep_bounded_encoder_state_beyond_ten_chunks(
|
||||
pipeline(_request(sampling=sampling))
|
||||
|
||||
assert pipeline.vae.encode_inputs == []
|
||||
assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 32
|
||||
# The stub encoder's history settles after its second block; the nine
|
||||
# blocks after that reuse the settled condition without an encode.
|
||||
assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 5
|
||||
assert all(torch.count_nonzero(video) == 0 for video in pipeline.vae.encoder.inputs[1:])
|
||||
assert pipeline._ar_sessions["world-1"].condition_fixed_point is not None
|
||||
encoder_cache = pipeline._ar_sessions["world-1"].encoder_cache
|
||||
assert sum(t.numel() * t.element_size() for t in encoder_cache) == pipeline._condition_encoder_cache_bytes()
|
||||
assert pipeline.vae._enc_feat_map == ["module-owned"]
|
||||
@@ -1943,6 +1970,76 @@ def test_condition_continues_real_causal_vae_beyond_initial_horizon(
|
||||
assert pipeline._ar_sessions == {}
|
||||
|
||||
|
||||
def test_condition_encoder_stops_at_its_fixed_point_and_restarts_with_the_session(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Once a block returns the history it was given, later blocks reuse the condition without encoding.
|
||||
|
||||
The whole-clip encode is the reference for every block, so reusing the
|
||||
stored condition is only allowed to be a speedup, never a change.
|
||||
"""
|
||||
pipeline = _tiny_condition_pipeline(monkeypatch)
|
||||
pipeline._ar_diffusion_kv_state = object()
|
||||
blocks = 12
|
||||
inputs = pipeline._parse_request(_request(sampling=_SamplingParams(num_frames=(blocks * 3 - 1) * 4 + 1)))
|
||||
with torch.inference_mode():
|
||||
expected = pipeline._prepare_condition(inputs, dtype=torch.float32)
|
||||
encoder_calls = []
|
||||
original_encoder = pipeline.vae.encoder.forward
|
||||
|
||||
def encoder(*args, **kwargs):
|
||||
encoder_calls.append(len(seen))
|
||||
return original_encoder(*args, **kwargs)
|
||||
|
||||
seen: list[torch.Tensor] = []
|
||||
|
||||
def generate_block(**kwargs):
|
||||
index = len(seen) % blocks
|
||||
torch.testing.assert_close(
|
||||
kwargs["condition"], expected[:, :, index * 3 : (index + 1) * 3], rtol=1e-6, atol=1e-6
|
||||
)
|
||||
seen.append(kwargs["condition"])
|
||||
return torch.zeros_like(kwargs["condition"][:, :16])
|
||||
|
||||
monkeypatch.setattr(pipeline.vae.encoder, "forward", encoder)
|
||||
monkeypatch.setattr(pipeline, "_generate_block", generate_block)
|
||||
|
||||
def tick(index):
|
||||
pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
|
||||
|
||||
settled_at = None
|
||||
for index in range(blocks):
|
||||
tick(index)
|
||||
state = pipeline._ar_sessions["world-1"]
|
||||
if settled_at is None and state.condition_fixed_point is not None:
|
||||
settled_at = index
|
||||
assert state.pending_encoder_cache is None and state.encoder_cache is not None
|
||||
assert settled_at is not None and 0 < settled_at < blocks - 1, settled_at
|
||||
# Three encoder passes per encoded block, none once the fixed point is stored.
|
||||
assert encoder_calls == [index for index in range(settled_at + 1) for _ in range(3)]
|
||||
# The stored condition is the session's own, not shared with the callers.
|
||||
assert all(seen[settled_at] is not later for later in seen[settled_at + 1 :])
|
||||
|
||||
# A session restarted from chunk 0 encodes again from the source image.
|
||||
pipeline.reset_ar_diffusion_session("world-1")
|
||||
tick(0)
|
||||
state = pipeline._ar_sessions["world-1"]
|
||||
assert state.condition_fixed_point is None
|
||||
assert encoder_calls[-3:] == [blocks] * 3
|
||||
|
||||
|
||||
def test_condition_encoder_fixed_point_is_released_with_the_session(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pipeline = _tiny_condition_pipeline(monkeypatch)
|
||||
pipeline._ar_diffusion_kv_state = object()
|
||||
monkeypatch.setattr(pipeline, "_generate_block", lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]))
|
||||
for index in range(8):
|
||||
pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
|
||||
state = pipeline._ar_sessions["world-1"]
|
||||
assert state.condition_fixed_point is not None
|
||||
pipeline.close_ar_diffusion_session("world-1")
|
||||
assert state.condition_fixed_point is None and state.encoder_cache is None
|
||||
|
||||
|
||||
def test_condition_encoder_interleaving_retry_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pipeline = _tiny_condition_pipeline(monkeypatch)
|
||||
pipeline._ar_diffusion_kv_state = object()
|
||||
@@ -2209,6 +2306,40 @@ def _stepwise_chunks(pipeline, state, ar_state):
|
||||
yield output
|
||||
|
||||
|
||||
def test_stepwise_block_shares_one_camera_cache_across_its_five_forwards() -> None:
|
||||
"""The four probes and the commit of one stepwise block reuse one camera cache; the next block gets its own.
|
||||
|
||||
Request mode creates one cache per loop; the stepwise path's forwards are
|
||||
separate scheduler steps, so the pipeline holds the block's cache and
|
||||
passes it to each one.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "latent"
|
||||
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
assert len(outputs) == 2 and len(transformer.calls) == 10
|
||||
caches = [call["camera_cache"] for call in transformer.calls]
|
||||
assert all(c is not None for c in caches)
|
||||
assert len({id(c) for c in caches[:5]}) == 1 and len({id(c) for c in caches[5:]}) == 1
|
||||
assert caches[0] is not caches[5]
|
||||
# The block's cache does not outlive its commit.
|
||||
assert "camera_cache" not in state.extra
|
||||
|
||||
|
||||
def test_pipeline_declares_chunk_step_grouping() -> None:
|
||||
"""A realtime block's probes and commit may run without a scheduler cycle between them."""
|
||||
module = _load_pipeline_module()
|
||||
assert module.LingBotWorldCausalDMDPipeline.supports_chunk_step_grouping is True
|
||||
assert supports_chunk_step_grouping(_pipeline(module))
|
||||
|
||||
|
||||
def test_pipeline_declares_step_execution_support() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
@@ -2337,7 +2468,10 @@ def test_stepwise_matches_tick_transformer_trace_and_latents(num_chunks) -> None
|
||||
torch.testing.assert_close(step_call["camera_hidden_states"], tick_call["camera_hidden_states"])
|
||||
assert stepwise_state.extra["condition"].shape[2] == 3
|
||||
assert pipeline.vae.encode_inputs == []
|
||||
assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (num_chunks * 3 - 1)) * 2
|
||||
# The stub encoder's history settles after its second block, so both
|
||||
# modes encode two blocks and reuse the settled condition after that.
|
||||
encoded_blocks = min(num_chunks, 2)
|
||||
assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (encoded_blocks * 3 - 1)) * 2
|
||||
assert stepwise_fake.commits == ["main"] * num_chunks
|
||||
assert len(stepwise_outputs) == num_chunks and stepwise_outputs[-1].finished
|
||||
assert pipeline.vae.decode_inputs == []
|
||||
@@ -2406,21 +2540,16 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None:
|
||||
assert pipeline.vae.decode_inputs == []
|
||||
|
||||
|
||||
def _capture_video_processor(monkeypatch) -> list[tuple[tuple[int, ...], str]]:
|
||||
"""Record what reaches postprocess_video, and name each chunk's output."""
|
||||
import diffusers.video_processor as video_processor_module
|
||||
def _capture_pixel_conversion(monkeypatch) -> list[tuple[int, ...]]:
|
||||
"""Record each decoded block reaching the uint8 conversion, and name its output."""
|
||||
module = _load_pipeline_module()
|
||||
processed: list[tuple[int, ...]] = []
|
||||
|
||||
processed: list[tuple[tuple[int, ...], str]] = []
|
||||
def uint8_frames(video):
|
||||
processed.append(tuple(video.shape))
|
||||
return f"frames-{len(processed)}"
|
||||
|
||||
class VideoProcessor:
|
||||
def __init__(self, *, vae_scale_factor):
|
||||
assert vae_scale_factor == 8
|
||||
|
||||
def postprocess_video(self, video, *, output_type):
|
||||
processed.append((tuple(video.shape), output_type))
|
||||
return f"frames-{len(processed)}"
|
||||
|
||||
monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
|
||||
monkeypatch.setattr(module, "_uint8_frames", uint8_frames)
|
||||
return processed
|
||||
|
||||
|
||||
@@ -2434,7 +2563,7 @@ def _streaming_pipeline(module):
|
||||
def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
|
||||
"""A non-latent request must stream pixels under the primary "video" key."""
|
||||
module = _load_pipeline_module()
|
||||
processed = _capture_video_processor(monkeypatch)
|
||||
processed = _capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2443,7 +2572,7 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
assert [output_type for _, output_type in processed] == ["np", "np"]
|
||||
assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
|
||||
assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
|
||||
# Identity metadata survives the switch to pixel output.
|
||||
assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
|
||||
@@ -2457,7 +2586,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
|
||||
the same latents would: 9 frames, then 12.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
processed = _capture_video_processor(monkeypatch)
|
||||
processed = _capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2466,7 +2595,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
|
||||
_run_stepwise(pipeline, state)
|
||||
|
||||
assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
|
||||
assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
|
||||
# One opening frame for the session, not one per chunk.
|
||||
assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
|
||||
# The whole-clip decode path is not involved, and the module-owned cache
|
||||
@@ -2511,7 +2640,7 @@ def test_prepare_next_chunk_rejects_live_camera_with_scripted_paths() -> None:
|
||||
def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) -> None:
|
||||
"""Camera timelines must track the streaming decoder's 9-then-12 media counts."""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2536,7 +2665,7 @@ def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) ->
|
||||
def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
|
||||
"""Decoder state is keyed by request id and released with the AR session."""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(request_id="req-stream", num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2560,7 +2689,7 @@ def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
|
||||
def test_streaming_decode_state_is_dropped_on_session_reset(monkeypatch) -> None:
|
||||
"""A reset session must not resume the temporal context it just abandoned."""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(request_id="req-reset", num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2584,7 +2713,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
|
||||
interleaving the scheduler actually produces.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
processed = _capture_video_processor(monkeypatch)
|
||||
processed = _capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
states = []
|
||||
for request_id in ("req-a", "req-b"):
|
||||
@@ -2617,7 +2746,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
|
||||
assert pipeline.vae.decoder.first_chunk_flags == (
|
||||
[True, False, False] + [True, False, False] + [False] * 3 + [False] * 3
|
||||
)
|
||||
assert [shape for shape, _ in processed] == [
|
||||
assert processed == [
|
||||
(1, 3, 9, 16, 16),
|
||||
(1, 3, 9, 16, 16),
|
||||
(1, 3, 12, 16, 16),
|
||||
@@ -2628,7 +2757,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
|
||||
def test_streaming_decode_failure_drops_the_half_advanced_cache(monkeypatch) -> None:
|
||||
"""A failed chunk must not leave a cache a later chunk would continue from."""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(request_id="req-boom", num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2662,7 +2791,7 @@ def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None:
|
||||
never reaches, so this asserts on what ``post_quant_conv`` received.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
@@ -2695,7 +2824,7 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
|
||||
perf regression in the streaming decoder would be too.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
_capture_video_processor(monkeypatch)
|
||||
_capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _pipeline(
|
||||
module,
|
||||
transformer=_RecordingTransformer(),
|
||||
@@ -2723,21 +2852,18 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
|
||||
def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> None:
|
||||
"""Tiled decode drives its own cache, so a tiling shape keeps per-chunk decode."""
|
||||
module = _load_pipeline_module()
|
||||
processed = _capture_video_processor(monkeypatch)
|
||||
processed = _capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
# A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile.
|
||||
_enable_vae_tiling(pipeline, tile_sample_min=8)
|
||||
outputs = [
|
||||
pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), output_type="np", session_id="tiled")
|
||||
for _ in range(2)
|
||||
]
|
||||
outputs = [pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), session_id="tiled") for _ in range(2)]
|
||||
|
||||
# One whole-clip decode per AR block, each seeing only that block's frames.
|
||||
assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
|
||||
assert pipeline.vae.decoder.first_chunk_flags == []
|
||||
assert pipeline._streaming_decode_states == {}
|
||||
# Nine frames per chunk instead of 9 then 12: the frame loss being removed.
|
||||
assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
|
||||
assert processed == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
|
||||
assert outputs == ["frames-1", "frames-2"]
|
||||
|
||||
|
||||
@@ -2749,7 +2875,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
|
||||
on a configuration that would never have tiled.
|
||||
"""
|
||||
module = _load_pipeline_module()
|
||||
processed = _capture_video_processor(monkeypatch)
|
||||
processed = _capture_pixel_conversion(monkeypatch)
|
||||
pipeline = _streaming_pipeline(module)
|
||||
# A 2x2 latent is within a tile this size, so vae.decode would not tile.
|
||||
_enable_vae_tiling(pipeline, tile_sample_min=256)
|
||||
@@ -2761,7 +2887,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
|
||||
|
||||
assert pipeline.vae.decode_inputs == []
|
||||
assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
|
||||
assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
|
||||
assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
|
||||
|
||||
|
||||
def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
|
||||
@@ -2904,3 +3030,193 @@ def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
|
||||
|
||||
assert encoded == ["move through the room"]
|
||||
assert states[0].prompt_embeds is states[1].prompt_embeds
|
||||
|
||||
|
||||
def test_session_text_caches_publish_key_only_after_a_successful_build() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
state = module._LingBotARSessionState()
|
||||
invalidations: list[bool] = []
|
||||
fail_next = {"raise": False}
|
||||
|
||||
def fake_text_caches(prompt_embeds, *, invalidate):
|
||||
invalidations.append(invalidate)
|
||||
if fail_next["raise"]:
|
||||
fail_next["raise"] = False
|
||||
raise RuntimeError("injected failure inside _ar_text_caches")
|
||||
return ["caches"]
|
||||
|
||||
pipeline._ar_text_caches = fake_text_caches
|
||||
embeds = torch.ones(1, 512, 8)
|
||||
caches = lambda prompt, length=512, e=embeds: pipeline._session_text_caches( # noqa: E731
|
||||
e, prompt=prompt, max_sequence_length=length, session_state=state
|
||||
)
|
||||
|
||||
# First build of a session: runner state is empty -> no invalidation, key published.
|
||||
caches("A")
|
||||
assert invalidations == [False] and state.text_cache_key == ("A", 512, str(embeds.dtype))
|
||||
assert state.text_cache_dirty is False
|
||||
# Same inputs: reuse without invalidation.
|
||||
caches("A")
|
||||
assert invalidations == [False, False]
|
||||
# Regression (Codex review): B's build fails after invalidating -> no key is
|
||||
# left behind, so retrying A must invalidate again instead of consuming
|
||||
# partial/mixed K/V.
|
||||
fail_next["raise"] = True
|
||||
with pytest.raises(RuntimeError, match="injected failure"):
|
||||
caches("B")
|
||||
assert state.text_cache_key is None and state.text_cache_dirty is True
|
||||
caches("A")
|
||||
assert invalidations[-2:] == [True, True] and state.text_cache_key == ("A", 512, str(embeds.dtype))
|
||||
assert state.text_cache_dirty is False
|
||||
# The key covers every embedding-defining input: same text, other length or dtype -> rebuild.
|
||||
caches("A", length=256)
|
||||
caches("A", e=embeds.to(torch.bfloat16))
|
||||
assert invalidations[-2:] == [True, True]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# a source image is decoded once while the file behind it is unchanged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_png(path, colour=(10, 20, 30)):
|
||||
Image.new("RGB", (64, 64), colour).save(path)
|
||||
return path
|
||||
|
||||
|
||||
def _counting_decode(module):
|
||||
"""Replace the file decode with a counter, so the tests measure decodes rather than pixels."""
|
||||
decoded: list[str] = []
|
||||
original = module._decode_source_image_file
|
||||
|
||||
def decode(path):
|
||||
decoded.append(os.fspath(path))
|
||||
return Image.new("RGB", (64, 64), (1, 2, 3))
|
||||
|
||||
module._decode_source_image_file = decode
|
||||
return decoded, original
|
||||
|
||||
|
||||
def test_the_same_source_image_is_decoded_once(tmp_path: Path) -> None:
|
||||
module = _load_pipeline_module()
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
path = _write_png(tmp_path / "frame.png")
|
||||
decoded, original = _counting_decode(module)
|
||||
try:
|
||||
first = module._load_source_image(path)
|
||||
second = module._load_source_image(path)
|
||||
assert decoded == [os.fspath(os.path.realpath(path))] or len(decoded) == 1
|
||||
# Every caller owns its copy: mutating one must not reach the next caller or the cache.
|
||||
assert first is not second
|
||||
first.putpixel((0, 0), (255, 255, 255))
|
||||
assert module._load_source_image(path).getpixel((0, 0)) == (1, 2, 3)
|
||||
assert len(decoded) == 1
|
||||
finally:
|
||||
module._decode_source_image_file = original
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
|
||||
|
||||
def test_a_rewritten_source_image_is_decoded_again(tmp_path: Path) -> None:
|
||||
module = _load_pipeline_module()
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
path = _write_png(tmp_path / "frame.png")
|
||||
decoded, original = _counting_decode(module)
|
||||
try:
|
||||
module._load_source_image(path)
|
||||
module._load_source_image(path)
|
||||
assert len(decoded) == 1
|
||||
# Same path, different bytes: the mtime/size in the key must force a fresh decode.
|
||||
os.utime(path, ns=(0, 0))
|
||||
_write_png(path, colour=(200, 100, 50))
|
||||
module._load_source_image(path)
|
||||
assert len(decoded) == 2
|
||||
finally:
|
||||
module._decode_source_image_file = original
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
|
||||
|
||||
def test_an_image_rewritten_during_the_decode_is_never_cached(tmp_path: Path) -> None:
|
||||
module = _load_pipeline_module()
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
path = _write_png(tmp_path / "frame.png")
|
||||
decoded: list[str] = []
|
||||
original = module._decode_source_image_file
|
||||
|
||||
def decode_then_rewrite(p):
|
||||
decoded.append(os.fspath(p))
|
||||
_write_png(path, colour=(len(decoded), 0, 0))
|
||||
os.utime(path, ns=(len(decoded) * 1_000_000, len(decoded) * 1_000_000))
|
||||
return Image.new("RGB", (64, 64), (1, 2, 3))
|
||||
|
||||
module._decode_source_image_file = decode_then_rewrite
|
||||
try:
|
||||
module._load_source_image(path)
|
||||
# The entry is keyed by the fingerprint taken before the decode, which the rewrite invalidated: the
|
||||
# next call must decode again rather than serve the image that was decoded mid-rewrite.
|
||||
module._load_source_image(path)
|
||||
assert len(decoded) == 2
|
||||
finally:
|
||||
module._decode_source_image_file = original
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
|
||||
|
||||
def test_an_unstattable_path_keeps_upstream_validation(tmp_path: Path) -> None:
|
||||
module = _load_pipeline_module()
|
||||
module._decode_source_image_fingerprinted.cache_clear()
|
||||
missing = tmp_path / "does-not-exist.png"
|
||||
# The uncached path owns the error contract; caching must not change which exception a caller sees.
|
||||
with pytest.raises(ValueError):
|
||||
module._load_source_image(missing)
|
||||
assert module._decode_source_image_fingerprinted.cache_info().currsize == 0
|
||||
|
||||
|
||||
def test_the_text_cache_is_sized_by_cross_attention_heads_not_the_self_attention_share() -> None:
|
||||
"""Cross-attention replicates heads per rank, so its pool must not inherit the self-attention head count."""
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1)
|
||||
pipeline.transformer.blocks[0].cross_attn = SimpleNamespace(num_local_heads=4)
|
||||
|
||||
spec = pipeline.ar_diffusion_kv_cache_spec()
|
||||
|
||||
assert spec.num_kv_heads == 1
|
||||
assert spec.cross_attention_kv_heads == {"text": 4}
|
||||
assert spec.cross_attention_lengths == {"text": 512}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_type", ["np", "pil"])
|
||||
def test_streamed_chunk_is_uint8_frames_whatever_pixel_output_type_was_asked(monkeypatch, output_type) -> None:
|
||||
"""A realtime tick delivers (F, H, W, 3) uint8 for every pixel output type; only latent differs."""
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _streaming_pipeline(module)
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = output_type
|
||||
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
videos = [output.output["payload"]["video"] for output in outputs]
|
||||
assert [type(video) for video in videos] == [np.ndarray, np.ndarray]
|
||||
assert [video.dtype for video in videos] == [np.uint8, np.uint8]
|
||||
assert [video.shape for video in videos] == [(9, 16, 16, 3), (12, 16, 16, 3)]
|
||||
|
||||
|
||||
def test_uint8_frames_match_the_pil_conversion_byte_for_byte() -> None:
|
||||
"""The device-side conversion is the PIL path's arithmetic, so the bytes are the same."""
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
module = _load_pipeline_module()
|
||||
generator = torch.Generator().manual_seed(7)
|
||||
# Cover the clamp on both sides and values that sit exactly on rounding edges.
|
||||
video = torch.rand(1, 3, 5, 8, 8, generator=generator) * 2.4 - 1.2
|
||||
video[0, :, 0, 0, :4] = torch.tensor([-1.0, 1.0, 0.0, 1 / 255 - 1.0])
|
||||
|
||||
reference = VideoProcessor(vae_scale_factor=8).postprocess_video(video, output_type="pil")[0]
|
||||
expected = np.stack([np.asarray(frame) for frame in reference])
|
||||
|
||||
frames = module._uint8_frames(video)
|
||||
|
||||
assert frames.dtype == np.uint8 and frames.shape == (5, 8, 8, 3)
|
||||
assert frames.flags["C_CONTIGUOUS"]
|
||||
np.testing.assert_array_equal(frames, expected)
|
||||
@@ -103,6 +103,78 @@ class _StepPipeline:
|
||||
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
|
||||
|
||||
|
||||
class _ChunkedStepPipeline(_StepPipeline):
|
||||
"""Streaming stub: two chunks of two denoise steps each, one decode per chunk."""
|
||||
|
||||
supports_chunk_step_grouping = False
|
||||
|
||||
def prepare_encode(self, state, **kwargs):
|
||||
del kwargs
|
||||
self.prepare_calls += 1
|
||||
state.timesteps = [torch.tensor(10), torch.tensor(5), torch.tensor(10), torch.tensor(5)]
|
||||
state.latents = torch.tensor([0.0])
|
||||
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
|
||||
state.chunk_num_steps = 2
|
||||
state.total_chunks = 2
|
||||
state.chunk_index = 0
|
||||
state.step_in_chunk = 0
|
||||
return state
|
||||
|
||||
def step_scheduler(self, state, noise_pred, **kwargs):
|
||||
del noise_pred, kwargs
|
||||
self.scheduler_calls += 1
|
||||
state.step_in_chunk += 1
|
||||
state.step_index = state.step_in_chunk
|
||||
|
||||
def post_decode(self, state, **kwargs):
|
||||
del kwargs
|
||||
self.decode_calls += 1
|
||||
result = DiffusionOutput(output=torch.tensor([float(state.chunk_index)]))
|
||||
state.chunk_index += 1
|
||||
state.step_in_chunk = 0
|
||||
return result
|
||||
|
||||
|
||||
class _BatchDependentChunkPipeline(_ChunkedStepPipeline):
|
||||
"""Streaming stub whose noise depends on the batch's latents and timesteps, so a stale batch shows."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.seen: list[tuple[float, float]] = []
|
||||
|
||||
def prepare_encode(self, state, **kwargs):
|
||||
super().prepare_encode(state, **kwargs)
|
||||
state.latents = torch.tensor([1.0])
|
||||
return state
|
||||
|
||||
def denoise_step(self, input_batch, **kwargs):
|
||||
self.denoise_calls += 1
|
||||
latent = float(input_batch.latents.reshape(-1)[0])
|
||||
timestep = float(input_batch.timesteps.reshape(-1)[0])
|
||||
self.seen.append((latent, timestep))
|
||||
return torch.full((input_batch.latents.shape[0], 1), latent + timestep)
|
||||
|
||||
def step_scheduler(self, state, noise_pred, **kwargs):
|
||||
del kwargs
|
||||
self.scheduler_calls += 1
|
||||
state.latents = state.latents + noise_pred.reshape(-1)[0]
|
||||
state.step_in_chunk += 1
|
||||
state.step_index = state.step_in_chunk
|
||||
|
||||
def post_decode(self, state, **kwargs):
|
||||
self.final_latent = state.latents.reshape(-1)[0].clone()
|
||||
return super().post_decode(state, **kwargs)
|
||||
|
||||
|
||||
class _FailsMidChunkPipeline(_ChunkedStepPipeline):
|
||||
"""Streaming stub whose second denoise step of the first chunk raises."""
|
||||
|
||||
def step_scheduler(self, state, noise_pred, **kwargs):
|
||||
super().step_scheduler(state, noise_pred, **kwargs)
|
||||
if state.chunk_index == 0 and state.step_in_chunk == 2:
|
||||
raise RuntimeError("boom mid chunk")
|
||||
|
||||
|
||||
class _PerRequestErrorStepPipeline(_StepPipeline):
|
||||
def prepare_encode(self, state, **kwargs):
|
||||
if state.prompt == "fail-prepare":
|
||||
@@ -327,8 +399,9 @@ def _make_vllm_config():
|
||||
|
||||
def _make_runner(
|
||||
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
runner_cls: type[DiffusionModelRunner] = DiffusionModelRunner,
|
||||
):
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner = object.__new__(runner_cls)
|
||||
runner.vllm_config = _make_vllm_config()
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=None,
|
||||
@@ -511,6 +584,37 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
|
||||
torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_make_batch_identity_mapping_never_reads_the_device_tensor_back(monkeypatch):
|
||||
"""The default index mapping is built on the host; reading it back from the device would sync every step."""
|
||||
|
||||
def forbidden(self, *args, **kwargs):
|
||||
raise AssertionError("make_batch read a tensor back to the host")
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "tolist", forbidden)
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
|
||||
states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
|
||||
|
||||
batch = InputBatch.make_batch(states)
|
||||
repacked = InputBatch.make_batch(states, cached_batch=batch)
|
||||
|
||||
assert repacked is batch
|
||||
assert list(batch.idx_mapping_np) == [0, 1]
|
||||
assert batch.idx_mapping.dtype is torch.int32 and batch.idx_mapping.shape == (2,)
|
||||
assert batch.request_ids == ["req-1", "req-2"]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_make_batch_explicit_mapping_still_selects_and_orders_states():
|
||||
states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
|
||||
|
||||
batch = InputBatch.make_batch(states, idx_mapping=torch.tensor([1, 0]))
|
||||
|
||||
assert batch.request_ids == ["req-2", "req-1"]
|
||||
assert list(batch.idx_mapping_np) == [1, 0]
|
||||
torch.testing.assert_close(batch.latents, torch.tensor([[2.0], [1.0]]))
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_step_profiler_reports_denoise_step_as_diffuse():
|
||||
pipeline = _AutoDenoiseProfilerPipeline()
|
||||
@@ -628,6 +732,146 @@ class TestRunner:
|
||||
assert runner.pipeline.scheduler_calls == 2
|
||||
assert runner.pipeline.decode_calls == 1
|
||||
|
||||
@staticmethod
|
||||
def _make_grouping_runner(pipeline, *, grouping: bool):
|
||||
"""An AR runner without KV (so it inherits the stepwise entry) over a streaming stub pipeline."""
|
||||
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
||||
|
||||
runner = _make_runner(runner_cls=ARDiffusionModelRunner)
|
||||
runner.od_config.streaming_output = True
|
||||
runner.kv_cache = None
|
||||
runner.pipeline = pipeline
|
||||
runner.pipeline.supports_chunk_step_grouping = grouping
|
||||
return runner
|
||||
|
||||
@staticmethod
|
||||
def _drive_to_completion(runner, num_steps=4):
|
||||
outputs = [runner.execute_stepwise(_make_scheduler_output(_make_step_request(num_steps)))]
|
||||
step = 1
|
||||
while not outputs[-1].get_request_output("req-1").finished:
|
||||
outputs.append(runner.execute_stepwise(_make_cached_scheduler_output(step_id=step)))
|
||||
step += 1
|
||||
return [o.get_request_output("req-1") for o in outputs]
|
||||
|
||||
@pytest.mark.parametrize("grouping", [False, True])
|
||||
def test_ar_runner_runs_every_step_of_a_chunk_in_one_call_when_the_pipeline_declares_it(
|
||||
self, monkeypatch, grouping
|
||||
):
|
||||
"""With the capability declared, one AR runner call drives a whole chunk; the outputs are the same either way."""
|
||||
runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=grouping)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
denoise_calls_after_each_call = []
|
||||
original = runner.execute_stepwise
|
||||
|
||||
def counting(scheduler_output):
|
||||
output = original(scheduler_output)
|
||||
denoise_calls_after_each_call.append(runner.pipeline.denoise_calls)
|
||||
return output
|
||||
|
||||
runner.execute_stepwise = counting
|
||||
|
||||
per_call = self._drive_to_completion(runner)
|
||||
|
||||
# Chunk outputs, and the steps they were emitted at, are identical.
|
||||
emitted = [(o.step_index, float(o.result.output[0])) for o in per_call if o.result is not None]
|
||||
assert emitted == [(2, 0.0), (2, 1.0)]
|
||||
assert per_call[-1].finished is True
|
||||
assert runner.pipeline.denoise_calls == 4 and runner.pipeline.scheduler_calls == 4
|
||||
assert runner.pipeline.decode_calls == 2
|
||||
assert runner.pipeline.prepare_calls == 1
|
||||
# One call per chunk with the capability, one per step without; a
|
||||
# grouped call never crosses into the next chunk.
|
||||
assert denoise_calls_after_each_call == ([2, 4] if grouping else [1, 2, 3, 4])
|
||||
assert "req-1" not in runner.state_cache
|
||||
|
||||
def test_ar_runner_continues_a_grouped_chunk_as_a_cached_request(self, monkeypatch):
|
||||
"""The first step admits the new request; every further step of the chunk is a cached-request continuation."""
|
||||
runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=True)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
seen = []
|
||||
original_core = DiffusionModelRunner._execute_stepwise_core
|
||||
|
||||
def recording_core(self_, scheduler_output, **kwargs):
|
||||
seen.append(scheduler_output)
|
||||
return original_core(self_, scheduler_output, **kwargs)
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "_execute_stepwise_core", recording_core)
|
||||
|
||||
first = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4), finished_req_ids={"old"}))
|
||||
|
||||
assert first.get_request_output("req-1").result is not None
|
||||
assert runner.pipeline.prepare_calls == 1
|
||||
assert len(seen) == 2
|
||||
assert [n.request_id for n in seen[0].scheduled_new_reqs] == ["req-1"]
|
||||
continuation = seen[1]
|
||||
assert continuation.scheduled_new_reqs == []
|
||||
assert continuation.scheduled_cached_reqs.request_ids == ["req-1"]
|
||||
assert continuation.finished_req_ids == set()
|
||||
assert continuation.kv_prefetch_job is None and continuation.kv_connector_metadata is None
|
||||
assert continuation.step_id == seen[0].step_id
|
||||
|
||||
def test_ar_runner_grouped_steps_see_the_same_inputs_as_per_step_calls(self, monkeypatch):
|
||||
"""Every grouped step sees the latents and timestep step_scheduler just advanced, exactly as per-step calls do."""
|
||||
finals = {}
|
||||
seen = {}
|
||||
for grouping in (False, True):
|
||||
runner = self._make_grouping_runner(_BatchDependentChunkPipeline(), grouping=grouping)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
self._drive_to_completion(runner)
|
||||
finals[grouping] = float(runner.pipeline.final_latent)
|
||||
seen[grouping] = runner.pipeline.seen
|
||||
|
||||
# 1 -> +(1+10)=12 -> +(12+5)=29 -> +(29+10)=68 -> +(68+5)=141 in both modes.
|
||||
assert seen[False] == [(1.0, 10.0), (12.0, 5.0), (29.0, 10.0), (68.0, 5.0)]
|
||||
assert seen[True] == seen[False]
|
||||
assert finals[True] == finals[False] == 141.0
|
||||
|
||||
def test_ar_runner_grouped_chunk_stops_at_a_failing_step_and_drops_the_state(self, monkeypatch):
|
||||
"""A failure mid-chunk ends the grouped call with that step's error; nothing runs after it."""
|
||||
runner = self._make_grouping_runner(_FailsMidChunkPipeline(), grouping=True)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
|
||||
|
||||
failed = output.get_request_output("req-1")
|
||||
assert failed.finished is True
|
||||
assert failed.result is not None and "boom mid chunk" in failed.result.error
|
||||
assert runner.pipeline.denoise_calls == 2 and runner.pipeline.decode_calls == 0
|
||||
assert "req-1" not in runner.state_cache
|
||||
|
||||
def test_ar_runner_fails_a_chunk_that_never_advances(self, monkeypatch):
|
||||
pipeline = _ChunkedStepPipeline()
|
||||
monkeypatch.setattr(pipeline, "step_scheduler", lambda *args, **kwargs: None)
|
||||
runner = self._make_grouping_runner(pipeline, grouping=True)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
|
||||
failed = output.get_request_output("req-1")
|
||||
assert failed.finished
|
||||
assert "within 2 denoise steps" in failed.result.error
|
||||
assert pipeline.denoise_calls == 2
|
||||
assert "req-1" not in runner.state_cache
|
||||
|
||||
def test_base_runner_never_groups_and_the_ar_policy_needs_capability_streaming_and_one_request(self, monkeypatch):
|
||||
"""The shared runner is one step per call; the AR runner groups only a lone streaming request that declares it."""
|
||||
base = _make_runner()
|
||||
base.od_config.streaming_output = True
|
||||
base.pipeline = _ChunkedStepPipeline()
|
||||
base.pipeline.supports_chunk_step_grouping = True
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
first = DiffusionModelRunner.execute_stepwise(base, _make_scheduler_output(_make_step_request(4)))
|
||||
assert first.get_request_output("req-1").result is None and base.pipeline.denoise_calls == 1
|
||||
|
||||
ar = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=False)
|
||||
one = _make_cached_scheduler_output()
|
||||
two = _make_batch_scheduler_output([_make_step_request(4), _make_step_request(4)])
|
||||
two.scheduled_new_reqs[1].request_id = two.scheduled_new_reqs[1].req.request_id = "req-2"
|
||||
assert ar._groups_chunk_steps(one) is False # no capability declared
|
||||
ar.pipeline.supports_chunk_step_grouping = True
|
||||
assert ar._groups_chunk_steps(one) is True
|
||||
assert ar._groups_chunk_steps(two) is False # not for a batch
|
||||
ar.od_config.streaming_output = False
|
||||
assert ar._groups_chunk_steps(one) is False # not without streaming output
|
||||
|
||||
def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
runner.pipeline = _ProfilingStepPipeline()
|
||||
|
||||
@@ -64,6 +64,7 @@ _PUBLIC_OMNI = (
|
||||
"SPEAKER_MAX_UPLOADED",
|
||||
"SPEAKER_SAMPLES_DIR",
|
||||
"VLLM_OMNI_ABORT_TIMEOUT",
|
||||
"VLLM_OMNI_AR_DIFFUSION_KV_GATHER",
|
||||
"VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT",
|
||||
"VLLM_OMNI_EVENT_DRIVEN_ORCH",
|
||||
"VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS",
|
||||
|
||||
@@ -505,22 +505,8 @@ class WanDistCausalConv3d(nn.Conv3d):
|
||||
self._trim_cache: dict[int, tuple[int, int, int]] = {}
|
||||
|
||||
def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
|
||||
padding = list(self._padding)
|
||||
if cache_x is not None and padding[4] > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
|
||||
x = F.pad(x, padding)
|
||||
x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
|
||||
x,
|
||||
group=self.group,
|
||||
halo_size=self.halo_size,
|
||||
split_dim=self.split_dim,
|
||||
recv_top_buf=self._halo_recv_top_buf,
|
||||
recv_bottom_buf=self._halo_recv_bottom_buf,
|
||||
)
|
||||
shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
|
||||
x_padded, local_extent = self._assemble_input(x, cache_x)
|
||||
shift, start, upper_bound = self._get_trim_params(local_extent)
|
||||
if shift:
|
||||
x_padded = _narrow_along_dim(
|
||||
x_padded,
|
||||
@@ -532,6 +518,83 @@ class WanDistCausalConv3d(nn.Conv3d):
|
||||
out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
|
||||
return _zero_invalid_extent(out, split_dim=self.split_dim)
|
||||
|
||||
def _assemble_input(self, x: torch.Tensor, cache_x: torch.Tensor | None) -> tuple[torch.Tensor, int]:
|
||||
"""Build the conv input -- cached frames, causal/spatial zero padding and the halo columns -- in one pass.
|
||||
|
||||
The straightforward version wrote the activation three times per call: ``torch.cat`` of the cached
|
||||
frames, ``F.pad`` for the causal and spatial zeros, then another ``torch.cat`` to attach the halos the
|
||||
neighbours sent. Every causal conv of the decoder runs this on every latent frame, so at 480x832 that
|
||||
was the largest single source of memory traffic in the decode. Here the final tensor is allocated once
|
||||
and the cached frames and the new frames are copied straight into their slots; only the padding slices
|
||||
are zeroed. The halo rows are exchanged through contiguous send and receive buffers (a P2P op needs a
|
||||
contiguous tensor and the halo slot of a 5D tensor is a strided view) and copied into their slots. The
|
||||
conv sees the same tensor as before, byte for byte.
|
||||
|
||||
Returns the assembled tensor and the extent along the split dimension the trim parameters are keyed by
|
||||
(the extent before halos, as ``halo_exchange`` saw it).
|
||||
"""
|
||||
pad_w0, pad_w1, pad_h0, pad_h1, pad_t0, pad_t1 = self._padding
|
||||
cache_frames = 0
|
||||
if cache_x is not None and pad_t0 > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
cache_frames = int(cache_x.shape[2])
|
||||
pad_t0 -= cache_frames
|
||||
assert pad_t0 >= 0, "Wan temporal cache exceeds the causal padding"
|
||||
rank, world_size = _rank_world(self.group)
|
||||
halo = self.halo_size if world_size > 1 else 0
|
||||
halo_h = halo if self.split_dim == "height" else 0
|
||||
halo_w = halo if self.split_dim == "width" else 0
|
||||
|
||||
batch, channels, frames, height, width = x.shape
|
||||
top, left = pad_h0 + halo_h, pad_w0 + halo_w
|
||||
total_frames = pad_t0 + cache_frames + frames + pad_t1
|
||||
total_height = top + height + pad_h1 + halo_h
|
||||
total_width = left + width + pad_w1 + halo_w
|
||||
buf = torch.empty((batch, channels, total_frames, total_height, total_width), dtype=x.dtype, device=x.device)
|
||||
# Zero only what the conv must read as zeros: the causal frames in
|
||||
# front (and any trailing ones) and the spatial padding rows/columns.
|
||||
# The halo columns are filled below.
|
||||
if pad_t0:
|
||||
buf[:, :, :pad_t0].zero_()
|
||||
if pad_t1:
|
||||
buf[:, :, total_frames - pad_t1 :].zero_()
|
||||
if pad_h0:
|
||||
buf[:, :, :, halo_h:top].zero_()
|
||||
if pad_h1:
|
||||
buf[:, :, :, total_height - halo_h - pad_h1 : total_height - halo_h].zero_()
|
||||
if pad_w0:
|
||||
buf[:, :, :, :, halo_w:left].zero_()
|
||||
if pad_w1:
|
||||
buf[:, :, :, :, total_width - halo_w - pad_w1 : total_width - halo_w].zero_()
|
||||
interior = buf[:, :, pad_t0 : pad_t0 + cache_frames + frames, top : top + height, left : left + width]
|
||||
if cache_frames:
|
||||
interior[:, :, :cache_frames].copy_(cache_x)
|
||||
interior[:, :, cache_frames:].copy_(x)
|
||||
else:
|
||||
interior.copy_(x)
|
||||
|
||||
dim = self.split_tensor_dim
|
||||
local_extent = buf.shape[dim] - 2 * halo
|
||||
if halo:
|
||||
# What halo_exchange sends: the padded tensor's first and last
|
||||
# ``halo`` rows along the split dimension, padding included.
|
||||
top_row_ref = _narrow_along_dim(buf, dim, halo, halo)
|
||||
bottom_row_ref = _narrow_along_dim(buf, dim, buf.shape[dim] - 2 * halo, halo)
|
||||
self._halo_recv_top_buf = _ensure_recv_buf(self._halo_recv_top_buf, top_row_ref)
|
||||
self._halo_recv_bottom_buf = _ensure_recv_buf(self._halo_recv_bottom_buf, bottom_row_ref)
|
||||
_halo_exchange_p2p(
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
group=self.group,
|
||||
top_row_ref=top_row_ref,
|
||||
bottom_row_ref=bottom_row_ref,
|
||||
recv_top_buf=self._halo_recv_top_buf,
|
||||
recv_bottom_buf=self._halo_recv_bottom_buf,
|
||||
)
|
||||
_narrow_along_dim(buf, dim, 0, halo).copy_(self._halo_recv_top_buf)
|
||||
_narrow_along_dim(buf, dim, buf.shape[dim] - halo, halo).copy_(self._halo_recv_bottom_buf)
|
||||
return buf, local_extent
|
||||
|
||||
def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
|
||||
trim_params = self._trim_cache.get(local_extent)
|
||||
if trim_params is None:
|
||||
|
||||
@@ -92,6 +92,13 @@ class SupportsStepExecution(Protocol):
|
||||
``prepare_encode()`` (one-time request setup), ``denoise_step()``
|
||||
(one denoise forward), ``step_scheduler()`` (one scheduler update),
|
||||
and ``post_decode()`` (final decode).
|
||||
|
||||
A pipeline may additionally set the optional class attribute
|
||||
``supports_chunk_step_grouping = True`` to declare that its request state
|
||||
may advance through every denoise step of its current chunk without a
|
||||
serving-scheduler cycle in between (a capability, not a policy: the runner
|
||||
that owns the scheduling decision chooses whether to group steps). It is
|
||||
optional so it does not become part of the runtime protocol check.
|
||||
"""
|
||||
|
||||
supports_step_execution: ClassVar[bool] = True
|
||||
|
||||
@@ -22,6 +22,7 @@ from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
|
||||
from vllm_omni.diffusion.models.lingbot_world.transformer import (
|
||||
CameraModulationCache,
|
||||
LingBotAttentionCache,
|
||||
LingBotTransformerCache,
|
||||
)
|
||||
@@ -109,11 +110,16 @@ class LingBotDMDBlockRunner:
|
||||
start_frame: int,
|
||||
timestep_value: float,
|
||||
step_index: int,
|
||||
camera_cache: CameraModulationCache | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Predict flow for one denoise step.
|
||||
|
||||
A probe never writes KV: only the clean x0 of a finished block may
|
||||
enter the cache, which is ``commit_block_kv``'s job.
|
||||
|
||||
``camera_cache`` is the block's camera-modulation cache when the caller
|
||||
holds it across separate calls (the stepwise path); ``generate_block``
|
||||
opens the window itself instead.
|
||||
"""
|
||||
if not self.enforce_eager:
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
@@ -130,6 +136,7 @@ class LingBotDMDBlockRunner:
|
||||
cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=False),
|
||||
start_frame=start_frame,
|
||||
update_cache=False,
|
||||
camera_modulation_cache=camera_cache,
|
||||
)
|
||||
if flow_prediction.shape != current_latents.shape:
|
||||
raise RuntimeError(
|
||||
@@ -172,11 +179,13 @@ class LingBotDMDBlockRunner:
|
||||
cache: LingBotTransformerCache | None,
|
||||
ar: ARBlockContext | None,
|
||||
start_frame: int,
|
||||
camera_cache: CameraModulationCache | None = None,
|
||||
) -> None:
|
||||
"""Write the finished block's clean x0 into KV and commit its pages.
|
||||
|
||||
The fifth transformer call of a block, deliberately not a denoise step:
|
||||
on the stepwise path it belongs to ``post_decode()``.
|
||||
on the stepwise path it belongs to ``post_decode()``. ``camera_cache``
|
||||
is as in ``probe_step``.
|
||||
"""
|
||||
# Commit K/V only for the final clean block, never for noisy probes.
|
||||
cache_input = torch.cat((latents.to(dtype=condition.dtype), condition), dim=1)
|
||||
@@ -190,6 +199,7 @@ class LingBotDMDBlockRunner:
|
||||
cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=True),
|
||||
start_frame=start_frame,
|
||||
update_cache=True,
|
||||
camera_modulation_cache=camera_cache,
|
||||
)
|
||||
if ar is not None:
|
||||
ar.state.commit_paged_context(ar.branch)
|
||||
@@ -210,6 +220,10 @@ class LingBotDMDBlockRunner:
|
||||
"""Request mode: all probes of one block, then its commit."""
|
||||
block_shape = (1, self.transformer.config.out_channels, *condition.shape[2:5])
|
||||
current_latents = randn_tensor(block_shape, generator=generator, device=self.device, dtype=torch.float32)
|
||||
# Every forward below -- each probe and the commit -- reads the same camera trajectory, so the camera
|
||||
# injector is built once per block: one cache for the whole loop, exactly as the stepwise path holds
|
||||
# one per block across its separate calls.
|
||||
camera_cache = CameraModulationCache()
|
||||
for step_index, (timestep_value, sigma) in enumerate(schedule):
|
||||
flow_prediction = self.probe_step(
|
||||
current_latents=current_latents,
|
||||
@@ -221,6 +235,7 @@ class LingBotDMDBlockRunner:
|
||||
start_frame=start_frame,
|
||||
timestep_value=timestep_value,
|
||||
step_index=step_index,
|
||||
camera_cache=camera_cache,
|
||||
)
|
||||
next_sigma = schedule[step_index + 1][1] if step_index + 1 < len(schedule) else None
|
||||
current_latents = self.apply_transition(
|
||||
@@ -235,5 +250,6 @@ class LingBotDMDBlockRunner:
|
||||
cache=cache,
|
||||
ar=ar,
|
||||
start_frame=start_frame,
|
||||
camera_cache=camera_cache,
|
||||
)
|
||||
return current_latents
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import math
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
@@ -54,6 +55,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import (
|
||||
)
|
||||
from vllm_omni.diffusion.models.lingbot_world.dmd_block import ARBlockContext, LingBotDMDBlockRunner
|
||||
from vllm_omni.diffusion.models.lingbot_world.transformer import (
|
||||
CameraModulationCache,
|
||||
CausalLingBotWorldTransformer3DModel,
|
||||
LingBotAttentionCache,
|
||||
LingBotTransformerCache,
|
||||
@@ -91,7 +93,6 @@ logger = init_logger(__name__)
|
||||
_STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
from tqdm.std import tqdm as TqdmProgressBar
|
||||
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
@@ -152,6 +153,9 @@ class _LingBotRequestInputs:
|
||||
width: int
|
||||
num_frames: int
|
||||
num_latent_frames: int
|
||||
# Request mode returns pixels through VideoProcessor in this type. Step execution (realtime sessions)
|
||||
# returns every pixel type as ``(F, H, W, 3)`` uint8 frames, the layout the streaming endpoint muxes; only
|
||||
# ``"latent"`` differs there. See ``_decode_chunk_to_pixels``.
|
||||
output_type: str
|
||||
max_sequence_length: int
|
||||
flow_shift: float
|
||||
@@ -163,10 +167,20 @@ class _LingBotARSessionState:
|
||||
"""Model-owned state; attention tensors remain runner-owned."""
|
||||
|
||||
next_chunk_index: int = 0
|
||||
prompt: str | None = None
|
||||
# The umT5 embedding itself is reused by ``encode_prompt``'s own cache.
|
||||
# ``text_cache_key`` does the same job for the runner-owned cross-attention
|
||||
# K/V built from that embedding: it holds the (prompt, max_sequence_length,
|
||||
# dtype) the K/V was built from and is only published after a successful
|
||||
# build; ``text_cache_dirty`` is raised for the duration of a build, so a
|
||||
# build that fails forces the next tick to invalidate whatever its prompt is.
|
||||
text_cache_key: tuple[str, int, str] | None = None
|
||||
text_cache_dirty: bool = False
|
||||
generator_state: torch.Tensor | None = None
|
||||
encoder_cache: list[torch.Tensor | None] | None = None
|
||||
pending_encoder_cache: list[torch.Tensor | None] | None = None
|
||||
# The condition the encoder settles on once its history has stopped
|
||||
# changing; see ``_next_condition_chunk``.
|
||||
condition_fixed_point: torch.Tensor | None = None
|
||||
camera_tail: CameraTrajectory | None = None
|
||||
camera_pitch: float = 0.0
|
||||
|
||||
@@ -284,7 +298,70 @@ def _decode_source_image(image: PIL.Image.Image) -> PIL.Image.Image:
|
||||
raise ValueError(_SOURCE_IMAGE_ERROR) from None
|
||||
|
||||
|
||||
def _source_image_fingerprint(path: str | os.PathLike[str]) -> tuple[str, int, int, int, int, int]:
|
||||
"""Identify the bytes behind ``path``: the resolved path plus the stat fields that change when it is rewritten."""
|
||||
resolved = os.path.realpath(os.fspath(path))
|
||||
stat = os.stat(resolved)
|
||||
return resolved, stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns
|
||||
|
||||
|
||||
def _encoder_caches_identical(
|
||||
before: list[torch.Tensor | None] | None, after: list[torch.Tensor | None] | None
|
||||
) -> bool:
|
||||
"""Whether two encoder histories hold the same tensors bit for bit, slot by slot."""
|
||||
if before is None or after is None or len(before) != len(after):
|
||||
return False
|
||||
for x, y in zip(before, after, strict=True):
|
||||
if x is None or y is None:
|
||||
if x is not y:
|
||||
return False
|
||||
elif x.shape != y.shape or x.dtype != y.dtype or not torch.equal(x, y):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _uint8_frames(video: torch.Tensor) -> np.ndarray:
|
||||
"""Turn one decoded block ``(1, 3, F, H, W)`` in ``[-1, 1]`` into ``(F, H, W, 3)`` uint8 on the host.
|
||||
|
||||
This is the arithmetic ``VideoProcessor.postprocess_video`` applies for
|
||||
``"np"`` and ``"pil"`` -- denormalise, widen to float32, scale by 255 and
|
||||
round -- done on the device before the copy instead of on the host after
|
||||
it. The bytes are identical to the PIL path; what changes is that the
|
||||
frames cross to the host once, already in the layout and dtype the
|
||||
streaming endpoint muxes, instead of as a float32 copy that is then
|
||||
converted per frame while the next block waits.
|
||||
"""
|
||||
# Same order of operations as VideoProcessor: denormalise and clamp in
|
||||
# the decoder dtype, then widen, scale and round.
|
||||
frames = (video[0] / 2 + 0.5).clamp(0, 1).permute(1, 2, 3, 0)
|
||||
frames = frames.float().mul_(255).round_().to(torch.uint8)
|
||||
return np.ascontiguousarray(frames.cpu().numpy())
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _decode_source_image_fingerprinted(fingerprint: tuple[str, int, int, int, int, int]) -> PIL.Image.Image:
|
||||
"""The decoded image behind ``fingerprint`` (resolved path plus the stat fields that change on rewrite).
|
||||
|
||||
One entry: a request-mode caller that replays the same source path decodes it once. The realtime tick
|
||||
carries its image inline and never comes through here.
|
||||
"""
|
||||
return _decode_source_image_file(fingerprint[0])
|
||||
|
||||
|
||||
def _load_source_image(path: str | os.PathLike[str]) -> PIL.Image.Image:
|
||||
"""Decode the source image at ``path``, reusing the last decode while the file is unchanged.
|
||||
|
||||
Every caller gets its own copy, so a request that mutates its image cannot reach the next one. A path
|
||||
that cannot be stat'ed goes straight to the uncached decode, which owns the validation error contract.
|
||||
"""
|
||||
try:
|
||||
fingerprint = _source_image_fingerprint(path)
|
||||
except (OSError, TypeError, ValueError):
|
||||
return _decode_source_image_file(path)
|
||||
return _decode_source_image_fingerprinted(fingerprint).copy()
|
||||
|
||||
|
||||
def _decode_source_image_file(path: str | os.PathLike[str]) -> PIL.Image.Image:
|
||||
try:
|
||||
source_image = PIL.Image.open(path)
|
||||
except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError):
|
||||
@@ -483,6 +560,9 @@ class LingBotWorldCausalDMDPipeline(
|
||||
"""LingBot-World v2 I2V generation with a request-local causal cache."""
|
||||
|
||||
supports_step_execution: ClassVar[bool] = True
|
||||
# A realtime block is four probes and a commit over one prepared chunk
|
||||
# state; nothing between those steps needs the serving scheduler.
|
||||
supports_chunk_step_grouping: ClassVar[bool] = True
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
@@ -637,12 +717,21 @@ class LingBotWorldCausalDMDPipeline(
|
||||
raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.")
|
||||
condition_latent_frames = int(self.transformer.config.num_frames_per_block)
|
||||
condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels)
|
||||
element_size = torch.empty((), dtype=self.transformer.dtype).element_size()
|
||||
condition_bytes_per_session = (
|
||||
condition_channels
|
||||
* condition_latent_frames
|
||||
* latent_height
|
||||
* latent_width
|
||||
* torch.empty((), dtype=self.transformer.dtype).element_size()
|
||||
condition_channels * condition_latent_frames * latent_height * latent_width * element_size
|
||||
)
|
||||
# Each layer retains scale and shift for the rank-local camera tokens
|
||||
# across all forwards of a block, including between scheduler calls.
|
||||
camera_tokens = tokens_per_frame * condition_latent_frames
|
||||
sp_size = int(self.od_config.parallel_config.ulysses_degree)
|
||||
camera_bytes_per_session = (
|
||||
2
|
||||
* int(self.transformer.config.num_layers)
|
||||
* ((camera_tokens + sp_size - 1) // sp_size)
|
||||
* int(self.transformer.config.num_attention_heads)
|
||||
* int(self.transformer.config.attention_head_dim)
|
||||
* element_size
|
||||
)
|
||||
return ARDiffusionKVCacheSpec(
|
||||
num_layers=int(self.transformer.config.num_layers),
|
||||
@@ -658,10 +747,14 @@ class LingBotWorldCausalDMDPipeline(
|
||||
ARDiffusionCrossAttentionKVSpec(
|
||||
self._AR_TEXT_CACHE,
|
||||
_MAX_SEQUENCE_LENGTH,
|
||||
# Cross-attention keeps every local head on every rank rather than sharding heads across
|
||||
# the Ulysses group, so the text pool holds that many heads, not the self-attention share.
|
||||
num_kv_heads=int(self.transformer.blocks[0].cross_attn.num_local_heads),
|
||||
),
|
||||
),
|
||||
model_owned_state_bytes_per_session=(
|
||||
condition_bytes_per_session
|
||||
+ camera_bytes_per_session
|
||||
+ 2 * self._condition_encoder_cache_bytes() # Committed and in-flight histories.
|
||||
+ self._streaming_decode_bytes_per_session()
|
||||
),
|
||||
@@ -733,6 +826,7 @@ class LingBotWorldCausalDMDPipeline(
|
||||
if state is not None:
|
||||
state.encoder_cache = None
|
||||
state.pending_encoder_cache = None
|
||||
state.condition_fixed_point = None
|
||||
|
||||
def reset_ar_diffusion_session(self, session_id: str) -> None:
|
||||
self._release_condition_encoder_state(session_id)
|
||||
@@ -988,6 +1082,50 @@ class LingBotWorldCausalDMDPipeline(
|
||||
)
|
||||
return condition
|
||||
|
||||
def _next_condition_chunk(
|
||||
self,
|
||||
inputs: _LingBotRequestInputs,
|
||||
*,
|
||||
start_frame: int,
|
||||
session_state: _LingBotARSessionState,
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
|
||||
"""The session's condition for the block at ``start_frame``, encoding only while it still changes.
|
||||
|
||||
After a session's opening block the encoder is fed the same all-zero
|
||||
pixel frames on every block, so the only thing that can make one
|
||||
block's condition differ from the last is the encoder's own history.
|
||||
The history is a causal cache with finite temporal reach, and once the
|
||||
zero frames have flushed the source image out of it the encoder is at
|
||||
a fixed point: identical input and identical state produce the same
|
||||
output and the same state on every later block. That fixed point is
|
||||
detected, not assumed after some number of blocks: a block whose
|
||||
returned history is bit-identical to the history it was given is the
|
||||
proof, and from then on the stored condition is returned and the
|
||||
encoder is not run. Measured at 480x832 the encode is 178 ms of a
|
||||
block and the fixed point arrives by the eleventh; the served
|
||||
rollouts this was measured on (40 blocks) are then byte-identical.
|
||||
|
||||
The opening block resets the fixed point, so a session that restarts
|
||||
from a new source image encodes again until it settles.
|
||||
"""
|
||||
if start_frame == 0:
|
||||
session_state.condition_fixed_point = None
|
||||
elif session_state.condition_fixed_point is not None:
|
||||
assert session_state.encoder_cache is not None
|
||||
# The history is unchanged by construction; the pending copy is
|
||||
# what post_decode commits, exactly as an encoded block's would be.
|
||||
return session_state.condition_fixed_point.clone(), list(session_state.encoder_cache)
|
||||
condition, pending_cache = self._prepare_condition_chunk(
|
||||
inputs,
|
||||
start_frame=start_frame,
|
||||
encoder_cache=session_state.encoder_cache,
|
||||
dtype=dtype,
|
||||
)
|
||||
if start_frame > 0 and _encoder_caches_identical(session_state.encoder_cache, pending_cache):
|
||||
session_state.condition_fixed_point = condition.clone()
|
||||
return condition, pending_cache
|
||||
|
||||
@torch.no_grad()
|
||||
def _prepare_condition_chunk(
|
||||
self,
|
||||
@@ -1160,10 +1298,8 @@ class LingBotWorldCausalDMDPipeline(
|
||||
for block in self.transformer.blocks:
|
||||
cross_attention = block.cross_attn
|
||||
shape = (cross_attention.num_local_heads, cross_attention.head_dim)
|
||||
key = cross_attention.shard_kv_heads(
|
||||
cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape)
|
||||
)
|
||||
value = cross_attention.shard_kv_heads(cross_attention.v(projected_text).unflatten(2, shape))
|
||||
key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape)
|
||||
value = cross_attention.v(projected_text).unflatten(2, shape)
|
||||
yield key, value
|
||||
|
||||
state.populate_cross_attention(
|
||||
@@ -1243,6 +1379,34 @@ class LingBotWorldCausalDMDPipeline(
|
||||
progress_bar=progress_bar,
|
||||
)
|
||||
|
||||
def _session_text_caches(
|
||||
self,
|
||||
prompt_embeds: torch.Tensor,
|
||||
*,
|
||||
prompt: str,
|
||||
max_sequence_length: int,
|
||||
session_state: _LingBotARSessionState,
|
||||
) -> list[LingBotAttentionCache]:
|
||||
"""Runner-owned cross-attention K/V for this tick, rebuilt when its inputs change.
|
||||
|
||||
The K/V is keyed by everything that defines the prompt embedding
|
||||
(text, sequence length, dtype). The first build of a session needs no
|
||||
invalidation (the runner state is empty); a key change does. The dirty
|
||||
flag is raised for the duration of the build and cleared only after
|
||||
``_ar_text_caches`` returns, so a build that fails after invalidating
|
||||
or partially writing the caches makes the next tick, whatever its
|
||||
prompt, invalidate again instead of consuming partial K/V.
|
||||
"""
|
||||
text_key = (prompt, int(max_sequence_length), str(prompt_embeds.dtype))
|
||||
previous = session_state.text_cache_key
|
||||
invalidate = session_state.text_cache_dirty or (previous is not None and previous != text_key)
|
||||
session_state.text_cache_dirty = True
|
||||
session_state.text_cache_key = None
|
||||
caches = self._ar_text_caches(prompt_embeds, invalidate=invalidate)
|
||||
session_state.text_cache_key = text_key
|
||||
session_state.text_cache_dirty = False
|
||||
return caches
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str,
|
||||
@@ -1345,6 +1509,9 @@ class LingBotWorldCausalDMDPipeline(
|
||||
)
|
||||
dtype = self.transformer.dtype
|
||||
# Phase 1: turn all three user inputs into DiT-ready conditions.
|
||||
# ``encode_prompt`` reuses an already-encoded prompt itself; the
|
||||
# runner-owned cross-attention K/V is keyed the same way in
|
||||
# ``_session_text_caches`` below.
|
||||
prompt_embeds = self.encode_prompt(
|
||||
inputs.prompt,
|
||||
max_sequence_length=inputs.max_sequence_length,
|
||||
@@ -1354,10 +1521,10 @@ class LingBotWorldCausalDMDPipeline(
|
||||
condition = self._prepare_condition(inputs, dtype=dtype)
|
||||
else:
|
||||
assert session_state is not None
|
||||
condition, pending_encoder_cache = self._prepare_condition_chunk(
|
||||
condition, pending_encoder_cache = self._next_condition_chunk(
|
||||
inputs,
|
||||
start_frame=tick.chunk_index * block_frames,
|
||||
encoder_cache=session_state.encoder_cache,
|
||||
session_state=session_state,
|
||||
dtype=dtype,
|
||||
)
|
||||
camera_pitch: float | None = None
|
||||
@@ -1396,12 +1563,13 @@ class LingBotWorldCausalDMDPipeline(
|
||||
)
|
||||
else:
|
||||
assert session_state is not None
|
||||
prompt_changed = session_state.prompt is not None and session_state.prompt != inputs.prompt
|
||||
if session_state.generator_state is not None:
|
||||
inputs.generator.set_state(session_state.generator_state)
|
||||
ar_cross_attention = self._ar_text_caches(
|
||||
ar_cross_attention = self._session_text_caches(
|
||||
prompt_embeds,
|
||||
invalidate=prompt_changed,
|
||||
prompt=inputs.prompt,
|
||||
max_sequence_length=inputs.max_sequence_length,
|
||||
session_state=session_state,
|
||||
)
|
||||
generated_blocks: list[torch.Tensor] = []
|
||||
total_steps = (inputs.num_latent_frames // block_frames) * len(LINGBOT_DMD_TIMESTEPS)
|
||||
@@ -1437,7 +1605,6 @@ class LingBotWorldCausalDMDPipeline(
|
||||
cache = None
|
||||
if tick is not None:
|
||||
assert session_state is not None
|
||||
session_state.prompt = inputs.prompt
|
||||
session_state.generator_state = inputs.generator.get_state()
|
||||
session_state.encoder_cache = pending_encoder_cache
|
||||
session_state.camera_tail = camera_tail
|
||||
@@ -1536,11 +1703,15 @@ class LingBotWorldCausalDMDPipeline(
|
||||
if state is not None:
|
||||
state.release()
|
||||
|
||||
def _decode_chunk_to_pixels(
|
||||
self, latents: torch.Tensor, *, output_type: str, session_id: str | None = None
|
||||
) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]:
|
||||
def _decode_chunk_to_pixels(self, latents: torch.Tensor, *, session_id: str | None = None) -> np.ndarray:
|
||||
"""Decode one AR block so streaming consumers receive pixels, not latents.
|
||||
|
||||
The block comes back as ``(F, H, W, 3)`` uint8 whatever pixel output
|
||||
type the request named: the realtime endpoint muxes uint8 frames, so
|
||||
``"pil"`` and ``"np"`` only ever differed in how much work stood
|
||||
between the decoder and those bytes. Only ``"latent"`` means
|
||||
something else, and the caller keeps latents before reaching here.
|
||||
|
||||
With a ``session_id`` the block is decoded through that session's own
|
||||
temporal cache, so chunk *N + 1* continues chunk *N*: the causal decoder
|
||||
expands the session's opening latent frame to one raw frame and every
|
||||
@@ -1566,7 +1737,7 @@ class LingBotWorldCausalDMDPipeline(
|
||||
video = self.vae.decode(vae_latents, return_dict=False)[0]
|
||||
else:
|
||||
video = self._streaming_decode_chunk(decoder, vae_latents, cast(str, session_id))
|
||||
return self._video_processor().postprocess_video(video, output_type=output_type)
|
||||
return _uint8_frames(video)
|
||||
|
||||
def _streaming_decode_chunk(
|
||||
self, decoder: WanStreamingDecoder, vae_latents: torch.Tensor, session_id: str
|
||||
@@ -1587,15 +1758,6 @@ class LingBotWorldCausalDMDPipeline(
|
||||
self._release_streaming_decode_state(session_id)
|
||||
raise
|
||||
|
||||
def _video_processor(self) -> VideoProcessor:
|
||||
processor = getattr(self, "_cached_video_processor", None)
|
||||
if processor is None:
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
|
||||
self._cached_video_processor = processor
|
||||
return processor
|
||||
|
||||
def _require_bound_ar_state(self) -> None:
|
||||
if self._ar_diffusion_kv_state is None:
|
||||
raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.")
|
||||
@@ -1683,10 +1845,10 @@ class LingBotWorldCausalDMDPipeline(
|
||||
session_state = self._ar_sessions[state.request_id]
|
||||
if session_state.next_chunk_index != state.chunk_index:
|
||||
raise ValueError("LingBot condition chunks must be contiguous within the session.")
|
||||
condition, pending_encoder_cache = self._prepare_condition_chunk(
|
||||
condition, pending_encoder_cache = self._next_condition_chunk(
|
||||
inputs,
|
||||
start_frame=start_frame,
|
||||
encoder_cache=session_state.encoder_cache,
|
||||
session_state=session_state,
|
||||
dtype=extra["dtype"],
|
||||
)
|
||||
previous = extra.get("camera_tail")
|
||||
@@ -1780,6 +1942,12 @@ class LingBotWorldCausalDMDPipeline(
|
||||
raise ValueError("LingBot World requires the runner-provided torch.Generator.")
|
||||
extra["condition"] = condition
|
||||
extra["camera"] = camera
|
||||
# One camera-modulation cache per block: the four probes and the
|
||||
# commit below are separate scheduler steps, so the block runner
|
||||
# cannot open the reuse window around them the way generate_block
|
||||
# does around its own loop. Holding the cache here makes the window
|
||||
# span exactly this block, per request.
|
||||
extra["camera_cache"] = CameraModulationCache()
|
||||
extra["camera_tail"] = camera_tail
|
||||
extra["start_frame"] = start_frame
|
||||
extra["ar_cross_attention"] = self._ar_text_caches(
|
||||
@@ -1838,6 +2006,7 @@ class LingBotWorldCausalDMDPipeline(
|
||||
start_frame=int(extra["start_frame"]),
|
||||
timestep_value=float(schedule[step_in_chunk][0]),
|
||||
step_index=step_in_chunk,
|
||||
camera_cache=extra["camera_cache"],
|
||||
)
|
||||
|
||||
def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
|
||||
@@ -1880,6 +2049,7 @@ class LingBotWorldCausalDMDPipeline(
|
||||
cache=None,
|
||||
ar=self._ar_block_context(extra["ar_cross_attention"]),
|
||||
start_frame=int(extra["start_frame"]),
|
||||
camera_cache=extra.pop("camera_cache"),
|
||||
)
|
||||
completed_chunk_index = state.chunk_index
|
||||
inputs: _LingBotRequestInputs = extra["inputs"]
|
||||
@@ -1889,13 +2059,7 @@ class LingBotWorldCausalDMDPipeline(
|
||||
if inputs.output_type == "latent":
|
||||
payload: dict[str, Any] = {"latents": latents}
|
||||
else:
|
||||
payload = {
|
||||
"video": self._decode_chunk_to_pixels(
|
||||
latents,
|
||||
output_type=inputs.output_type,
|
||||
session_id=state.request_id,
|
||||
)
|
||||
}
|
||||
payload = {"video": self._decode_chunk_to_pixels(latents, session_id=state.request_id)}
|
||||
output = {
|
||||
"payload": payload,
|
||||
"metadata": {
|
||||
|
||||
@@ -24,7 +24,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
|
||||
from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D, all_to_all_5D
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
@@ -329,9 +329,18 @@ class LingBotSelfAttention(nn.Module):
|
||||
key = self.rotary_embedding(key, cos, sin)
|
||||
|
||||
if self.ulysses_world_size > 1:
|
||||
query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
|
||||
key = SeqAllToAll4D.apply(self.ulysses_group, key, 2, 1, False)
|
||||
value = SeqAllToAll4D.apply(self.ulysses_group, value, 2, 1, False)
|
||||
# One fused sequence->head all-to-all for q/k/v instead of three:
|
||||
# (B, S/N, 3, H, D) -> (B, S, 3, H/N, D). Same head-group assignment
|
||||
# as the per-tensor 4D exchange (rank r receives head group r), so
|
||||
# the result is a pure permutation of the three-call version, with a
|
||||
# 3x larger message and one third of the latency-bound collectives.
|
||||
qkv = all_to_all_5D(
|
||||
torch.stack((query, key, value), dim=2),
|
||||
scatter_idx=3,
|
||||
gather_idx=1,
|
||||
group=self.ulysses_group,
|
||||
)
|
||||
query, key, value = qkv.unbind(2)
|
||||
|
||||
if isinstance(cache, ARDiffusionPagedLayerInputs):
|
||||
if query.shape[0] != 1:
|
||||
@@ -421,13 +430,11 @@ class LingBotCrossAttention(nn.Module):
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.num_local_heads = num_heads // tp_size
|
||||
self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state()
|
||||
if self.num_local_heads % self.ulysses_world_size:
|
||||
raise ValueError(
|
||||
"LingBot local cross-attention heads must be divisible by the Ulysses degree: "
|
||||
f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}."
|
||||
)
|
||||
self.num_sp_heads = self.num_local_heads // self.ulysses_world_size
|
||||
# Cross-attention keeps every local head on every Ulysses rank and attends the rank's own sequence
|
||||
# shard. Each query token attends the whole text sequence independently of the others, so there is
|
||||
# nothing to exchange: the sequence split Ulysses already gives us is the only split needed, and both
|
||||
# all-to-all collectives -- one for the query, one for the output -- disappear. The cost is that the text K/V is
|
||||
# replicated per rank instead of head-sharded; ar_diffusion_kv_cache_spec sizes the pool for it.
|
||||
self.tp_inner_dim = self.num_local_heads * self.head_dim
|
||||
|
||||
self.q = ColumnParallelLinear(
|
||||
@@ -465,9 +472,9 @@ class LingBotCrossAttention(nn.Module):
|
||||
self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
||||
self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_sp_heads,
|
||||
num_heads=self.num_local_heads,
|
||||
head_size=self.head_dim,
|
||||
num_kv_heads=self.num_sp_heads,
|
||||
num_kv_heads=self.num_local_heads,
|
||||
softmax_scale=self.head_dim**-0.5,
|
||||
causal=False,
|
||||
role="cross",
|
||||
@@ -477,14 +484,6 @@ class LingBotCrossAttention(nn.Module):
|
||||
disable_kv_quant=True,
|
||||
)
|
||||
|
||||
def shard_kv_heads(self, value: torch.Tensor) -> torch.Tensor:
|
||||
"""Select this Ulysses rank's projected K/V heads."""
|
||||
|
||||
if self.ulysses_world_size == 1:
|
||||
return value
|
||||
start = self.ulysses_rank * self.num_sp_heads
|
||||
return value[:, :, start : start + self.num_sp_heads].clone(memory_format=torch.contiguous_format)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -494,8 +493,6 @@ class LingBotCrossAttention(nn.Module):
|
||||
) -> tuple[torch.Tensor, LingBotAttentionCache]:
|
||||
query = self.norm_q(self.q(hidden_states))
|
||||
query = query.unflatten(2, (self.num_local_heads, self.head_dim))
|
||||
if self.ulysses_world_size > 1:
|
||||
query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
|
||||
|
||||
# Text K/V is constant within a request and is projected once per layer.
|
||||
if cache is None:
|
||||
@@ -503,8 +500,8 @@ class LingBotCrossAttention(nn.Module):
|
||||
raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.")
|
||||
key = self.norm_k(self.k(encoder_hidden_states))
|
||||
value = self.v(encoder_hidden_states)
|
||||
key = self.shard_kv_heads(key.unflatten(2, (self.num_local_heads, self.head_dim)))
|
||||
value = self.shard_kv_heads(value.unflatten(2, (self.num_local_heads, self.head_dim)))
|
||||
key = key.unflatten(2, (self.num_local_heads, self.head_dim))
|
||||
value = value.unflatten(2, (self.num_local_heads, self.head_dim))
|
||||
cache = LingBotAttentionCache(
|
||||
key=key,
|
||||
value=value,
|
||||
@@ -517,8 +514,6 @@ class LingBotCrossAttention(nn.Module):
|
||||
value = cache.value[:, : cache.end]
|
||||
|
||||
output = self.attn(query, key, value)
|
||||
if self.ulysses_world_size > 1:
|
||||
output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False)
|
||||
return self.o(output.flatten(2, 3)), cache
|
||||
|
||||
|
||||
@@ -592,6 +587,16 @@ class LingBotAttentionBlock(nn.Module):
|
||||
self.cam_scale_layer = nn.Linear(dim, dim)
|
||||
self.cam_shift_layer = nn.Linear(dim, dim)
|
||||
|
||||
def camera_modulation(self, camera_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build this block's camera scale and shift from the patchified camera tokens.
|
||||
|
||||
Depends only on ``camera_hidden_states``, so a caller holding the camera trajectory fixed -- every
|
||||
denoise forward of one AR chunk does -- can build this once and pass it back through ``forward``.
|
||||
"""
|
||||
camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
|
||||
camera_features = camera_features + camera_hidden_states
|
||||
return self.cam_scale_layer(camera_features), self.cam_shift_layer(camera_features)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -605,6 +610,7 @@ class LingBotAttentionBlock(nn.Module):
|
||||
sink_tokens: int,
|
||||
update_cache: bool,
|
||||
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
||||
camera_modulation: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, LingBotAttentionCache]:
|
||||
batch_size, token_count, dim = hidden_states.shape
|
||||
if (
|
||||
@@ -638,10 +644,9 @@ class LingBotAttentionBlock(nn.Module):
|
||||
hidden_grid = hidden_grid + attention_output.unflatten(1, (num_frames, tokens_per_frame)) * gate_msa
|
||||
hidden_states = hidden_grid.flatten(1, 2).to(hidden_states.dtype)
|
||||
|
||||
camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
|
||||
camera_features = camera_features + camera_hidden_states
|
||||
camera_scale = self.cam_scale_layer(camera_features)
|
||||
camera_shift = self.cam_shift_layer(camera_features)
|
||||
camera_scale, camera_shift = (
|
||||
self.camera_modulation(camera_hidden_states) if camera_modulation is None else camera_modulation
|
||||
)
|
||||
hidden_states = ((1 + camera_scale) * hidden_states + camera_shift).to(hidden_states.dtype)
|
||||
|
||||
attention_output, cross_cache = self.cross_attn(
|
||||
@@ -763,6 +768,22 @@ def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Ten
|
||||
return phase.cos().float(), phase.sin().float()
|
||||
|
||||
|
||||
class CameraModulationCache:
|
||||
"""Per-block camera scale and shift, held by the caller across the forwards of one AR block.
|
||||
|
||||
The camera injector is four projections and a SiLU per block and reads only the camera tokens, which do
|
||||
not change between the denoise steps and the clean commit of one block. The caller that owns that
|
||||
invariant creates one of these per block and passes it to every forward of the block; ``forward`` fills a
|
||||
block's entry on first use and reads it after. Its lifetime is the caller's: the AR runner keeps one per
|
||||
block, ``generate_block`` one per loop, and nothing here checks the camera tensor for change.
|
||||
"""
|
||||
|
||||
__slots__ = ("entries",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.entries: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
|
||||
|
||||
|
||||
class CausalLingBotWorldTransformer3DModel(nn.Module):
|
||||
"""Checkpoint-compatible causal LingBot World video transformer."""
|
||||
|
||||
@@ -922,6 +943,27 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
|
||||
cosine, sine = _rope_axis(rope_max_seq_len, axis_dim)
|
||||
self.register_buffer(f"_rope_{axis}_cosine", cosine, persistent=False)
|
||||
self.register_buffer(f"_rope_{axis}_sine", sine, persistent=False)
|
||||
# One static-address buffer per shape for the per-token timestep
|
||||
# projection; see _stage_timestep_projection.
|
||||
self._timestep_projection_buffers: dict[tuple[tuple[int, ...], torch.dtype, torch.device], torch.Tensor] = {}
|
||||
|
||||
def _stage_timestep_projection(self, projection: torch.Tensor) -> torch.Tensor:
|
||||
"""Hand the blocks the timestep projection from a buffer whose address never changes.
|
||||
|
||||
Every regional block graph takes the projection as an input. Under CUDA-graph replay an input that
|
||||
arrives in a fresh tensor each forward is copied into the graph's placeholder on every replay: at
|
||||
480x832 with four Ulysses ranks that is a 72 MB copy per block, forty times per forward, two hundred
|
||||
times per AR chunk. Copying the projection once per forward into a buffer marked static lets every
|
||||
replay read it in place. The values the blocks see are the same; only where they live changes.
|
||||
"""
|
||||
key = (tuple(projection.shape), projection.dtype, projection.device)
|
||||
buffer = self._timestep_projection_buffers.get(key)
|
||||
if buffer is None:
|
||||
buffer = torch.empty(projection.shape, dtype=projection.dtype, device=projection.device)
|
||||
torch._dynamo.mark_static_address(buffer)
|
||||
self._timestep_projection_buffers[key] = buffer
|
||||
buffer.copy_(projection)
|
||||
return buffer
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
@@ -1178,6 +1220,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
|
||||
cache: LingBotTransformerCache,
|
||||
start_frame: int,
|
||||
update_cache: bool,
|
||||
camera_modulation_cache: CameraModulationCache | None = None,
|
||||
) -> torch.Tensor:
|
||||
if not isinstance(update_cache, bool):
|
||||
raise ValueError("update_cache must be a boolean.")
|
||||
@@ -1272,16 +1315,25 @@ class CausalLingBotWorldTransformer3DModel(nn.Module):
|
||||
"LingBot SP input hooks did not shard all conditioning tensors consistently; "
|
||||
f"expected {local_tokens} tokens per rank. Check SP hook registration and input dimensions."
|
||||
)
|
||||
timestep_projection = self._stage_timestep_projection(timestep_projection)
|
||||
rotary_emb = (cosine, sine)
|
||||
# Phase 3: each layer receives its own cache entry. Text K/V is passed
|
||||
# only when absent; the returned cache is stored for subsequent DMD
|
||||
# steps and causal blocks in this request.
|
||||
# steps and causal blocks in this request. A caller-held camera cache
|
||||
# gives every block its modulation once per AR block.
|
||||
for index, block in enumerate(self.blocks):
|
||||
camera_modulation = None
|
||||
if camera_modulation_cache is not None:
|
||||
camera_modulation = camera_modulation_cache.entries.get(index)
|
||||
if camera_modulation is None:
|
||||
camera_modulation = block.camera_modulation(camera_hidden_states)
|
||||
camera_modulation_cache.entries[index] = camera_modulation
|
||||
hidden_states, cross_cache = block(
|
||||
hidden_states,
|
||||
projected_text if cache.cross_attention[index] is None else None,
|
||||
timestep_projection,
|
||||
camera_hidden_states,
|
||||
camera_modulation=camera_modulation,
|
||||
self_cache=cache.self_attention[index],
|
||||
cross_cache=cache.cross_attention[index],
|
||||
current_start=current_start,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Diffusion input-batch structures following the MRV2-style vLLM layout.
|
||||
|
||||
Request states remain the only persistent source of truth. Static tensors are
|
||||
@@ -67,19 +67,26 @@ def _select_states(
|
||||
raise ValueError("Cannot build InputBatch from empty states.")
|
||||
|
||||
if idx_mapping is None:
|
||||
# The identity mapping is known on the host; build the device copy as
|
||||
# a launch and the host copy from NumPy rather than reading the device
|
||||
# tensor back. This runs on every step of every stepwise request, and
|
||||
# a device-to-host read here would stall the host until the queued
|
||||
# denoise work drained, so it must not be a synchronising operation.
|
||||
device = states[0].latents.device if states[0].latents is not None else None
|
||||
idx_mapping = torch.arange(len(states), dtype=torch.int32, device=device)
|
||||
idx_mapping_np = np.arange(len(states), dtype=np.int32)
|
||||
else:
|
||||
if idx_mapping.ndim != 1:
|
||||
raise ValueError("idx_mapping must be a 1D tensor.")
|
||||
idx_mapping = idx_mapping.to(dtype=torch.int32)
|
||||
idx_mapping_np = idx_mapping.detach().cpu().numpy()
|
||||
|
||||
selected_states: list[StepRequestState] = []
|
||||
for batch_idx, state_idx in enumerate(idx_mapping.tolist()):
|
||||
for batch_idx, state_idx in enumerate(idx_mapping_np.tolist()):
|
||||
if state_idx < 0 or state_idx >= len(states):
|
||||
raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
|
||||
selected_states.append(states[state_idx])
|
||||
return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
|
||||
return selected_states, idx_mapping, idx_mapping_np
|
||||
|
||||
|
||||
def _prepare_request_ids(states: Sequence[StepRequestState]) -> list[str]:
|
||||
|
||||
@@ -40,12 +40,18 @@ class ARDiffusionCrossAttentionKVSpec:
|
||||
|
||||
name: str
|
||||
num_tokens: int
|
||||
# Head count of this cross-attention cache. Cross-attention need not be sharded the same way as
|
||||
# self-attention: a model whose cross-attention keeps every local head on every rank stores that many heads
|
||||
# here instead of its self-attention share. None means "same as self-attention", the existing behaviour.
|
||||
num_kv_heads: int | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.name:
|
||||
raise ValueError("AR-Diffusion cross-attention cache names must be non-empty")
|
||||
if self.num_tokens <= 0:
|
||||
raise ValueError(f"AR-Diffusion cross-attention num_tokens must be positive, got {self.num_tokens}")
|
||||
if self.num_kv_heads is not None and self.num_kv_heads <= 0:
|
||||
raise ValueError(f"AR-Diffusion cross-attention num_kv_heads must be positive, got {self.num_kv_heads}")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -135,6 +141,11 @@ class ARDiffusionKVCacheSpec:
|
||||
def cross_attention_lengths(self) -> dict[str, int]:
|
||||
return {cache.name: cache.num_tokens for cache in self.cross_attention}
|
||||
|
||||
@property
|
||||
def cross_attention_kv_heads(self) -> dict[str, int]:
|
||||
"""Per-cache head counts, for the caches that do not use the self-attention head count."""
|
||||
return {cache.name: cache.num_kv_heads for cache in self.cross_attention if cache.num_kv_heads is not None}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsARDiffusionPipeline(Protocol):
|
||||
@@ -181,3 +192,12 @@ class SupportsARDiffusionWarmup(Protocol):
|
||||
def ar_diffusion_warmup_requests(self, session_id: str) -> Iterable[OmniDiffusionRequest]:
|
||||
"""Yield requests for compiled shapes, each carrying ``session_id``."""
|
||||
...
|
||||
|
||||
|
||||
def supports_chunk_step_grouping(pipeline: object) -> bool:
|
||||
"""Whether a step-execution ``pipeline`` declares that one chunk's steps may run without scheduler cycles.
|
||||
|
||||
An AR-runner policy input, consulted only after the pipeline has already been admitted to step
|
||||
execution; LingBot-World declares it as a class attribute.
|
||||
"""
|
||||
return bool(getattr(pipeline, "supports_chunk_step_grouping", False))
|
||||
@@ -1,10 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Configuration for AR-Diffusion engine-level KV cache management."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
KV_GATHER_ENV = "VLLM_OMNI_AR_DIFFUSION_KV_GATHER"
|
||||
|
||||
|
||||
def contiguous_kv_gather_enabled() -> bool:
|
||||
"""Whether the experimental contiguous-K/V gather attention path is on (``VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1``).
|
||||
|
||||
The one place the switch is read: the KV manager consults it when it budgets and allocates the history
|
||||
staging buffers, and the attention dispatch when it picks the path, so the two cannot disagree. It is the
|
||||
only consumer of ``ARDiffusionKVConfig.reuse_history_staging``.
|
||||
"""
|
||||
return os.environ.get(KV_GATHER_ENV, "0") == "1"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ARDiffusionKVConfig:
|
||||
@@ -30,6 +44,12 @@ class ARDiffusionKVConfig:
|
||||
# DiT graphs for every window-fill shape at load time via a synthetic rollout,
|
||||
# so the serving run is fast from the first chunk. No effect when eager.
|
||||
warmup_cudagraph: bool = True
|
||||
# Keep one contiguous K/V staging buffer per layer and refresh only the tokens that changed.
|
||||
# Within one AR block every forward attends the same history and differs only in the current
|
||||
# chunk, so re-gathering the whole visible window per forward re-copies bytes that did not move.
|
||||
# Off by default: it costs two max-sequence buffers per layer, per rank. See
|
||||
# ARDiffusionPagedForwardContext.history_staging.
|
||||
reuse_history_staging: bool = False
|
||||
# Also capture the post-window-boundary (reset-cycle) forward during warm-up.
|
||||
warmup_capture_reset: bool = False
|
||||
|
||||
|
||||
@@ -15,7 +15,8 @@ from __future__ import annotations
|
||||
import inspect
|
||||
import os
|
||||
from collections.abc import Collection, Iterable, Sequence
|
||||
from typing import cast
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
@@ -29,7 +30,11 @@ from vllm.v1.request import RequestStatus
|
||||
|
||||
from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor
|
||||
from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import (
|
||||
KV_GATHER_ENV,
|
||||
ARDiffusionKVConfig,
|
||||
contiguous_kv_gather_enabled,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
|
||||
ChunkWindowManager,
|
||||
ChunkWindowSpec,
|
||||
@@ -158,6 +163,18 @@ def build_kv_manager(
|
||||
return KVCacheManager(config, **kwargs)
|
||||
|
||||
|
||||
@dataclass
|
||||
class HistoryStagingState:
|
||||
"""What the staged history window currently holds: whose session and which visible blocks.
|
||||
|
||||
Owned by the cache next to the buffers themselves; ``ARDiffusionPagedForwardContext`` reads it to decide
|
||||
whether a forward may keep the staged history and writes back what it staged.
|
||||
"""
|
||||
|
||||
adapter: Any | None = None # weakref to the session adapter the window was staged for
|
||||
signature: tuple[Any, ...] | None = None
|
||||
|
||||
|
||||
class ARDiffusionKVCache:
|
||||
"""Own the paged KV pool and KV-branch-local storage for one model.
|
||||
|
||||
@@ -180,6 +197,7 @@ class ARDiffusionKVCache:
|
||||
kv_branches: tuple[ARDiffusionKVBranchSpec, ...],
|
||||
session_capacity: int,
|
||||
cross_attention_lengths: dict[str, int] | None = None,
|
||||
cross_attention_kv_heads: dict[str, int] | None = None,
|
||||
device: torch.device | None = None,
|
||||
frames_per_block: int = 1,
|
||||
max_scratch_tokens_per_branch: int = 0,
|
||||
@@ -225,6 +243,18 @@ class ARDiffusionKVCache:
|
||||
self.head_size = head_size
|
||||
self.dtype = dtype
|
||||
self.cross_attention_lengths = dict(cross_attention_lengths or {})
|
||||
# A cross-attention cache may hold more heads than the self-attention share: a model that keeps every
|
||||
# local head on every rank, rather than sharding heads across the sequence-parallel group, stores the
|
||||
# full local set here. Absent means the self-attention head count, which is what every caller did before.
|
||||
self.cross_attention_kv_heads = {
|
||||
name: int((cross_attention_kv_heads or {}).get(name, num_kv_heads)) for name in self.cross_attention_lengths
|
||||
}
|
||||
invalid_heads = {n: h for n, h in self.cross_attention_kv_heads.items() if h <= 0}
|
||||
if invalid_heads:
|
||||
raise ValueError(f"cross_attention_kv_heads must be positive, got {invalid_heads}")
|
||||
unknown_heads = set(cross_attention_kv_heads or {}) - set(self.cross_attention_lengths)
|
||||
if unknown_heads:
|
||||
raise ValueError(f"cross_attention_kv_heads names unknown caches: {sorted(unknown_heads)}")
|
||||
invalid_cross = {name: length for name, length in self.cross_attention_lengths.items() if length <= 0}
|
||||
if invalid_cross:
|
||||
raise ValueError(f"cross_attention_lengths must be positive, got {invalid_cross}")
|
||||
@@ -274,20 +304,46 @@ class ARDiffusionKVCache:
|
||||
self.scratch_reserved_bytes = self.scratch_num_blocks * page_size_bytes
|
||||
self.model_owned_state_bytes_per_session = model_owned_state_bytes_per_session
|
||||
|
||||
def _cross_pool_bytes(length: int) -> int:
|
||||
return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
|
||||
def _cross_pool_bytes(length: int, heads: int) -> int:
|
||||
return int(2 * len(self.kv_branches) * length * heads * head_size * dtype.itemsize * num_layers)
|
||||
|
||||
self.cross_attention_bytes_per_session = sum(
|
||||
_cross_pool_bytes(length) for length in self.cross_attention_lengths.values()
|
||||
_cross_pool_bytes(length, self.cross_attention_kv_heads[name])
|
||||
for name, length in self.cross_attention_lengths.items()
|
||||
)
|
||||
|
||||
def _required_managed_blocks(capacity: int) -> int:
|
||||
resident_per_session = config.sink_chunks + config.window_chunks
|
||||
return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2
|
||||
|
||||
# reuse_history_staging keeps one contiguous K and V buffer per layer,
|
||||
# sized to the padded visible window (sink + window plus the
|
||||
# action-capacity block build_block_table always reserves), for the
|
||||
# whole worker. It is allocated lazily on the first prepared forward,
|
||||
# which is after admission, so it has to be reserved here or admission
|
||||
# can succeed and that first forward run out of memory. Only the
|
||||
# contiguous-gather attention path consumes it.
|
||||
self.history_staging_reserved_bytes = 0
|
||||
self.history_staging_tokens = 0
|
||||
if config.reuse_history_staging:
|
||||
if contiguous_kv_gather_enabled():
|
||||
self.history_staging_tokens = (
|
||||
self.spec.sliding_window + config.sink_chunks * config.chunk_size + block_size
|
||||
)
|
||||
self.history_staging_reserved_bytes = int(
|
||||
2 * num_layers * self.history_staging_tokens * num_kv_heads * head_size * dtype.itemsize
|
||||
)
|
||||
else:
|
||||
_log.warning(
|
||||
"reuse_history_staging is set but the contiguous K/V gather path (%s=1) is off: "
|
||||
"staging has no consumer and is neither budgeted nor allocated.",
|
||||
KV_GATHER_ENV,
|
||||
)
|
||||
|
||||
def _required_bytes(capacity: int) -> int:
|
||||
return (
|
||||
self.scratch_reserved_bytes
|
||||
+ self.history_staging_reserved_bytes
|
||||
+ capacity * self.cross_attention_bytes_per_session
|
||||
+ capacity * self.model_owned_state_bytes_per_session
|
||||
+ _required_managed_blocks(capacity) * page_size_bytes
|
||||
@@ -298,7 +354,7 @@ class ARDiffusionKVCache:
|
||||
raise ValueError(
|
||||
"AR-Diffusion available device memory cannot fit one session: "
|
||||
f"available={available_bytes} bytes, required={one_session_bytes} bytes "
|
||||
"(managed self-attention + cross-attention + scratch + model-owned state)."
|
||||
"(managed self-attention + cross-attention + scratch + model-owned state + K/V staging)."
|
||||
)
|
||||
self.memory_budget_bytes = max(self.configured_memory_budget_bytes, one_session_bytes)
|
||||
if self.memory_budget_bytes > self.configured_memory_budget_bytes:
|
||||
@@ -325,6 +381,7 @@ class ARDiffusionKVCache:
|
||||
self_attn_budget_bytes = (
|
||||
self.memory_budget_bytes
|
||||
- self.scratch_reserved_bytes
|
||||
- self.history_staging_reserved_bytes
|
||||
- self.cross_attention_reserved_bytes
|
||||
- self.model_owned_state_reserved_bytes
|
||||
)
|
||||
@@ -374,6 +431,20 @@ class ARDiffusionKVCache:
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
# reuse_history_staging: one contiguous (K, V) pair per layer, sized to the padded visible window this
|
||||
# cache admits and budgeted above, owned here next to the pools it stages from. Allocated outside any
|
||||
# capture and marked static, so compiled regions treat it as a stable input rather than a new tensor.
|
||||
self.history_staging: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||||
self.history_staging_state = HistoryStagingState()
|
||||
if self.history_staging_tokens and device is not None:
|
||||
for _ in range(num_layers):
|
||||
pair = (
|
||||
torch.empty((self.history_staging_tokens, num_kv_heads, head_size), device=device, dtype=dtype),
|
||||
torch.empty((self.history_staging_tokens, num_kv_heads, head_size), device=device, dtype=dtype),
|
||||
)
|
||||
for buffer in pair:
|
||||
torch._dynamo.mark_static_address(buffer)
|
||||
self.history_staging.append(pair)
|
||||
|
||||
# -- cross-attention pool access -------------------------------------------
|
||||
# Cross-attn KV is static once populated — write once (from text encoder),
|
||||
@@ -455,7 +526,7 @@ class ARDiffusionKVCache:
|
||||
if not self._allocate_tensors:
|
||||
raise RuntimeError("AR-Diffusion cross-attention tensors require a configured pool device")
|
||||
|
||||
shape = (length, self.num_kv_heads, self.head_size)
|
||||
shape = (length, self.cross_attention_kv_heads[cache_name], self.head_size)
|
||||
expected_input_shape = (1, *shape)
|
||||
k_pool = [torch.empty(shape, dtype=self.dtype, device=self.device) for _ in range(self.num_layers)]
|
||||
v_pool = [torch.empty(shape, dtype=self.dtype, device=self.device) for _ in range(self.num_layers)]
|
||||
|
||||
@@ -115,11 +115,27 @@ def allocate_kv_pool_with_views(
|
||||
cache_shape = (num_blocks, block_size, num_kv_heads, head_dim)
|
||||
flat_shape = (num_blocks * block_size, num_kv_heads, head_dim)
|
||||
for _ in range(num_layers):
|
||||
k = torch.empty(cache_shape, dtype=dtype, device=device)
|
||||
v = torch.empty(cache_shape, dtype=dtype, device=device)
|
||||
kv_pools.append([k, v])
|
||||
k_pools.append(k.reshape(flat_shape))
|
||||
v_pools.append(v.reshape(flat_shape))
|
||||
# The flat slot views are what the paged-write custom op mutates inside the
|
||||
# compiled block graphs, so they must be the *base* allocations: dynamo only
|
||||
# honours mark_static_address on the tensor object it sees as a graph input,
|
||||
# and a view of a static base is still classified as a mutated input, which
|
||||
# makes inductor skip CUDA graphs for every block ("skipping cudagraphs due
|
||||
# to mutated inputs"). The block-table layout is therefore the view here.
|
||||
k = torch.empty(flat_shape, dtype=dtype, device=device)
|
||||
v = torch.empty(flat_shape, dtype=dtype, device=device)
|
||||
for pool in (k, v):
|
||||
# Block 0 is the manager's null block: block tables are tail-padded with
|
||||
# it and it is never written. The contiguous-K/V gather path may read
|
||||
# (masked) rows from it, so it must hold finite values; every other
|
||||
# block is fully written before it is read.
|
||||
pool[:block_size].zero_()
|
||||
# Session-lifetime storage: tell dynamo/inductor the address is static so
|
||||
# CUDA-graph trees (mode="reduce-overhead") mutate it in place instead of
|
||||
# copying ~300 MB per layer per replay (or skipping the graph).
|
||||
torch._dynamo.mark_static_address(pool)
|
||||
kv_pools.append([k.view(cache_shape), v.view(cache_shape)])
|
||||
k_pools.append(k)
|
||||
v_pools.append(v)
|
||||
return kv_pools, k_pools, v_pools
|
||||
|
||||
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Paged self-attention helpers for AR-Diffusion KV reuse."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import weakref
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, ClassVar, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import contiguous_kv_gather_enabled
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
|
||||
|
||||
_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
|
||||
@@ -21,6 +24,19 @@ def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
|
||||
return t
|
||||
|
||||
|
||||
def _to_device_async(t: torch.Tensor, device: torch.device) -> torch.Tensor:
|
||||
"""Host->device copy that does not stall the CPU.
|
||||
|
||||
``tensor.to(device)`` from pageable memory is synchronous: the CPU blocks until
|
||||
every kernel already queued on the stream has finished. Pinned + non_blocking
|
||||
keeps the CPU running ahead (the caching host allocator keeps the pinned
|
||||
source alive until the copy's stream event completes).
|
||||
"""
|
||||
if torch.device(device).type != "cuda" or t.device.type != "cpu":
|
||||
return t.to(device=device)
|
||||
return t.pin_memory().to(device=device, non_blocking=True)
|
||||
|
||||
|
||||
class ARDiffusionPagedLayerInputs(NamedTuple):
|
||||
"""Compiled-region payload for one layer's paged self-attention.
|
||||
|
||||
@@ -48,6 +64,14 @@ class ARDiffusionPagedLayerInputs(NamedTuple):
|
||||
seq_lens: torch.Tensor
|
||||
max_query_len: int
|
||||
max_seq_len: int
|
||||
# Contiguous K/V staging for this layer, when reuse_history_staging is on. ``reuse_history`` is a host
|
||||
# bool decided once per forward (``_prepare_history_staging``): it selects one of two compiled variants.
|
||||
stage_key: torch.Tensor | None = None
|
||||
stage_value: torch.Tensor | None = None
|
||||
reuse_history: bool = False
|
||||
# Index of the current chunk's first block within the staged window: right after the visible history,
|
||||
# which is not the end of the padded block table (that carries at least one action-capacity block).
|
||||
stage_first_block: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -104,15 +128,15 @@ class ARDiffusionPagedForwardContext:
|
||||
start_block = start // self.block_size
|
||||
self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]]
|
||||
positions = torch.arange(start, start + self.seq_len, dtype=torch.long)
|
||||
self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
|
||||
self.current_video_slot_mapping = _to_device_async(
|
||||
compute_slot_mapping(table, positions, self.block_size), device
|
||||
)
|
||||
else:
|
||||
self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, n_blocks)
|
||||
positions = torch.arange(self.seq_len, dtype=torch.long)
|
||||
self.current_video_slot_mapping = compute_slot_mapping(
|
||||
self.current_video_block_ids,
|
||||
positions,
|
||||
self.block_size,
|
||||
).to(device=device)
|
||||
self.current_video_slot_mapping = _to_device_async(
|
||||
compute_slot_mapping(self.current_video_block_ids, positions, self.block_size), device
|
||||
)
|
||||
self._allocated_video = True
|
||||
|
||||
def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
|
||||
@@ -135,11 +159,9 @@ class ARDiffusionPagedForwardContext:
|
||||
action_blocks,
|
||||
)
|
||||
positions = torch.arange(action_len, dtype=torch.long)
|
||||
self.action_slot_mapping = compute_slot_mapping(
|
||||
self.action_scratch_block_ids,
|
||||
positions,
|
||||
self.block_size,
|
||||
).to(device=device)
|
||||
self.action_slot_mapping = _to_device_async(
|
||||
compute_slot_mapping(self.action_scratch_block_ids, positions, self.block_size), device
|
||||
)
|
||||
self._action_len = action_len
|
||||
|
||||
def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
|
||||
@@ -192,9 +214,12 @@ class ARDiffusionPagedForwardContext:
|
||||
self.query_len = int(query_len)
|
||||
self.kv_len = int(video_len + action_len)
|
||||
max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
|
||||
block_table = torch.tensor([padded], dtype=torch.int32, device=device)
|
||||
query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
|
||||
seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
|
||||
# Built on the host once per AR block; the copies go through the same
|
||||
# pinned, non-blocking path as the slot mappings above so the CPU does
|
||||
# not wait for the stream to drain three times before the first layer.
|
||||
block_table = _to_device_async(torch.tensor([padded], dtype=torch.int32), device)
|
||||
query_start_loc = _to_device_async(torch.tensor([0, self.query_len], dtype=torch.int32), device)
|
||||
seq_lens = _to_device_async(torch.tensor([self.kv_len], dtype=torch.int32), device)
|
||||
return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len
|
||||
|
||||
def prepare(self, device: torch.device, action_len: int, query_len: int) -> None:
|
||||
@@ -218,15 +243,71 @@ class ARDiffusionPagedForwardContext:
|
||||
) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
|
||||
if self.action_slot_mapping is None:
|
||||
self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
|
||||
self._prepare_history_staging(action_len)
|
||||
self._prepared = True
|
||||
|
||||
def _prepare_history_staging(self, action_len: int) -> None:
|
||||
"""Decide whether this forward can keep the history already staged, and from where.
|
||||
|
||||
The staged window is only reusable when nothing about it moved: the same session adapter, the same
|
||||
visible history blocks, the same sequence bounds. Any difference -- a new AR block, a slid window, a
|
||||
different request -- and the buffer is rebuilt from the pools.
|
||||
|
||||
The decision is a host bool, deliberately. It has to be, for the gather to actually be skipped: a
|
||||
device flag could only select between two results already computed. It is the same value for all
|
||||
layers of a forward, so it costs two compiled variants in total rather than one per layer.
|
||||
"""
|
||||
self.staging_enabled = False
|
||||
self.reuse_history = False
|
||||
self.stage_first_block = 0
|
||||
# The manager allocates the staging pairs only when reuse_history_staging is on and the contiguous
|
||||
# gather path -- their only consumer -- is switched on.
|
||||
if not self.kv_cache.history_staging:
|
||||
return
|
||||
if action_len or self.block_table.shape[0] != 1:
|
||||
# Action tokens live in scratch blocks outside the video window; staging them is not modelled.
|
||||
return
|
||||
capacity = int(self.kv_cache.history_staging[0][0].shape[0])
|
||||
if int(self.max_seq_len) > capacity:
|
||||
raise RuntimeError(
|
||||
f"history staging buffers hold {capacity} tokens but this forward stages {int(self.max_seq_len)}"
|
||||
)
|
||||
state = self.kv_cache.history_staging_state
|
||||
signature = (
|
||||
int(self.adapter.num_computed_tokens),
|
||||
tuple(self.history_block_ids),
|
||||
int(self.max_seq_len),
|
||||
int(self.seq_len),
|
||||
)
|
||||
same_session = state.adapter is not None and state.adapter() is self.adapter
|
||||
reuse = same_session and state.signature == signature
|
||||
state.adapter = weakref.ref(self.adapter)
|
||||
state.signature = signature
|
||||
self.staging_enabled = True
|
||||
self.reuse_history = bool(reuse)
|
||||
# ``kv_len`` is the visible video window here (no action tokens on this path): the current chunk's
|
||||
# blocks are its last ``num_current_video_blocks`` entries, and the table's trailing padding blocks
|
||||
# come after them.
|
||||
self.stage_first_block = self.kv_len // self.block_size - self.num_current_video_blocks
|
||||
|
||||
def history_staging(self, layer_idx: int) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
"""This layer's staging pair, narrowed to this forward's staged window (a prefix of the manager's buffer)."""
|
||||
if not self.staging_enabled:
|
||||
return None, None
|
||||
key, value = self.kv_cache.history_staging[layer_idx]
|
||||
tokens = int(self.max_seq_len)
|
||||
return key[:tokens], value[:tokens]
|
||||
|
||||
def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
|
||||
if not getattr(self, "_prepared", False):
|
||||
raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
|
||||
key_pool = self.kv_cache._k_pools[layer_idx]
|
||||
value_pool = self.kv_cache._v_pools[layer_idx]
|
||||
stage_key, stage_value = self.history_staging(layer_idx)
|
||||
return ARDiffusionPagedLayerInputs(
|
||||
layer_idx=_layer_idx_tensor(layer_idx),
|
||||
key_pool=self.kv_cache._k_pools[layer_idx],
|
||||
value_pool=self.kv_cache._v_pools[layer_idx],
|
||||
key_pool=key_pool,
|
||||
value_pool=value_pool,
|
||||
block_size=int(self.kv_cache.block_size),
|
||||
seq_len=int(self.seq_len),
|
||||
video_slots=self.current_video_slot_mapping,
|
||||
@@ -236,6 +317,10 @@ class ARDiffusionPagedForwardContext:
|
||||
seq_lens=self.seq_lens,
|
||||
max_query_len=int(self.max_query_len),
|
||||
max_seq_len=int(self.max_seq_len),
|
||||
stage_key=stage_key,
|
||||
stage_value=stage_value,
|
||||
reuse_history=self.reuse_history,
|
||||
stage_first_block=int(self.stage_first_block),
|
||||
)
|
||||
|
||||
def mark_committed(self) -> None:
|
||||
@@ -342,6 +427,23 @@ def _reference_paged_attention(
|
||||
_FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {}
|
||||
|
||||
|
||||
def _stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, n_blocks, block_size, *, first_block):
|
||||
"""Gather visible blocks from ``first_block`` onward into the staging buffers.
|
||||
|
||||
``first_block`` is 0 for a full restage and the index of the current chunk's first block when the
|
||||
history is already staged. Skipping the leading blocks is the whole point: they were gathered by an
|
||||
earlier forward of the same AR block and cannot have changed since.
|
||||
"""
|
||||
ids = block_ids[first_block:]
|
||||
blocks = n_blocks - first_block
|
||||
stage_key.view(n_blocks, block_size, *key_cache.shape[2:])[first_block:].copy_(
|
||||
key_cache.index_select(0, ids).view(blocks, block_size, *key_cache.shape[2:])
|
||||
)
|
||||
stage_value.view(n_blocks, block_size, *value_cache.shape[2:])[first_block:].copy_(
|
||||
value_cache.index_select(0, ids).view(blocks, block_size, *value_cache.shape[2:])
|
||||
)
|
||||
|
||||
|
||||
def _resolve_fa_version(head_size: int) -> int:
|
||||
# get_flash_attn_version -> current_platform.get_device_capability() is not
|
||||
# dynamo-traceable, and the answer is fixed per head size for the process.
|
||||
@@ -385,6 +487,11 @@ def ar_diffusion_paged_attention(
|
||||
max_seq_len: int,
|
||||
softmax_scale: float,
|
||||
causal: bool = False,
|
||||
stage_key: torch.Tensor | None = None,
|
||||
stage_value: torch.Tensor | None = None,
|
||||
reuse_history: bool = False,
|
||||
current_tokens: int = 0,
|
||||
stage_first_block: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Run non-causal paged attention over a vLLM block table.
|
||||
|
||||
@@ -437,6 +544,82 @@ def ar_diffusion_paged_attention(
|
||||
softmax_scale=float(softmax_scale),
|
||||
causal=causal,
|
||||
)
|
||||
elif contiguous_kv_gather_enabled() and block_table.shape[0] == 1:
|
||||
# Experimental (VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1): FA3's paged-KV path with the
|
||||
# frame-sized block (1560 tokens, not a multiple of the FA3 K/V tile)
|
||||
# runs ~15-17% slower than the same kernel on contiguous K/V. Gather the
|
||||
# visible blocks into a contiguous buffer once per layer and attend
|
||||
# without a block table; ``seqused_k`` masks the tail-padding block.
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
|
||||
fa_version = _resolve_fa_version(query_flat.shape[-1])
|
||||
block_size = key_cache.shape[1]
|
||||
n_blocks = int(max_seq_len) // block_size
|
||||
if n_blocks * block_size != int(max_seq_len):
|
||||
raise ValueError("the contiguous K/V gather path requires max_seq_len to be block-aligned")
|
||||
block_ids = block_table[0, :n_blocks].to(torch.long)
|
||||
if stage_key is None or stage_value is None:
|
||||
# Fresh allocations on purpose: a module-level cached buffer that is first
|
||||
# allocated inside a CUDA-graph-trees warm-up run lives in the graph pool
|
||||
# untracked ("tensor(s) in the cudagraph pool not tracked as outputs").
|
||||
k_buf = key_cache.index_select(0, block_ids)
|
||||
v_buf = value_cache.index_select(0, block_ids)
|
||||
k_flat = k_buf.view(n_blocks * block_size, *key_cache.shape[2:])
|
||||
v_flat = v_buf.view(n_blocks * block_size, *value_cache.shape[2:])
|
||||
else:
|
||||
# Caller-owned staging, allocated outside any capture and marked static, so the buffer this
|
||||
# writes into is the same address every forward and the graph pool never owns it.
|
||||
#
|
||||
# Only the current chunk moved. The block table is the visible history followed by the
|
||||
# current chunk's blocks and then padding (at least the action-capacity block, plus unused
|
||||
# window capacity while the window is still growing), so the rows this forward changed start
|
||||
# at ``stage_first_block`` -- the caller's count of visible history blocks -- and NOT at
|
||||
# ``n_blocks - current_blocks``, which would refresh padding and leave the current K/V stale.
|
||||
# Everything before them was staged by an earlier forward of the same AR block and is still
|
||||
# byte-for-byte what a full gather would produce. When the history did move -- new block,
|
||||
# slid window, different session -- reuse_history is 0 and the whole window is re-gathered,
|
||||
# which is the same work the unstaged path always does.
|
||||
k_flat = stage_key
|
||||
v_flat = stage_value
|
||||
if reuse_history:
|
||||
# The caller's metadata must describe a block-aligned current chunk inside the staged
|
||||
# window; anything else is a wrong offset, not a reason to quietly restage everything.
|
||||
current_blocks = current_tokens // block_size
|
||||
assert current_tokens > 0 and current_blocks * block_size == current_tokens, (
|
||||
f"staged reuse needs a block-aligned current chunk, got {current_tokens} tokens"
|
||||
)
|
||||
assert 0 <= stage_first_block and stage_first_block + current_blocks <= n_blocks, (
|
||||
f"stage_first_block={stage_first_block} + {current_blocks} blocks "
|
||||
f"exceeds the {n_blocks}-block window"
|
||||
)
|
||||
_stage_window(
|
||||
k_flat,
|
||||
v_flat,
|
||||
key_cache,
|
||||
value_cache,
|
||||
block_ids,
|
||||
n_blocks,
|
||||
block_size,
|
||||
first_block=stage_first_block if reuse_history else 0,
|
||||
)
|
||||
out = torch.empty_like(query_flat)
|
||||
# Varlen K/V: cu_seqlens_k = [0, kv_len] built on device from seq_lens
|
||||
# (no host sync). Rows past kv_len (the tail-padding block) are never
|
||||
# read, exactly like the ROCm gather path above.
|
||||
cu_seqlens_k = torch.cat((seq_lens.new_zeros(1), seq_lens.to(torch.int32)))
|
||||
flash_attn_varlen_func(
|
||||
q=query_flat,
|
||||
k=k_flat,
|
||||
v=v_flat,
|
||||
out=out,
|
||||
cu_seqlens_q=query_start_loc,
|
||||
max_seqlen_q=int(max_query_len),
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_k=int(max_seq_len),
|
||||
softmax_scale=float(softmax_scale),
|
||||
causal=causal,
|
||||
fa_version=fa_version,
|
||||
)
|
||||
else:
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
|
||||
@@ -488,6 +671,10 @@ def _paged_write_attn_impl(
|
||||
max_query_len: int,
|
||||
max_seq_len: int,
|
||||
softmax_scale: float,
|
||||
stage_key: torch.Tensor | None = None,
|
||||
stage_value: torch.Tensor | None = None,
|
||||
reuse_history: bool = False,
|
||||
stage_first_block: int = 0,
|
||||
) -> torch.Tensor:
|
||||
key_pool[video_slots] = k_curr.to(key_pool.dtype)
|
||||
value_pool[video_slots] = v_curr.to(value_pool.dtype)
|
||||
@@ -507,16 +694,23 @@ def _paged_write_attn_impl(
|
||||
max_seq_len=max_seq_len,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=False,
|
||||
stage_key=stage_key,
|
||||
stage_value=stage_value,
|
||||
reuse_history=reuse_history,
|
||||
current_tokens=int(k_curr.shape[0]),
|
||||
stage_first_block=stage_first_block,
|
||||
)
|
||||
|
||||
|
||||
# hasattr guard keeps registration idempotent across test re-imports that pop
|
||||
# the module from sys.modules (same as sage_attn3.py).
|
||||
if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
|
||||
|
||||
# Keep staging arguments required even when their value is None. The dispatcher
|
||||
# strips trailing defaults, and older PyTorch mutation handlers index the
|
||||
# positional arguments without restoring them before bumping version counters.
|
||||
@torch.library.custom_op(
|
||||
"vllm_omni::ar_diffusion_paged_write_attn",
|
||||
mutates_args=("key_pool", "value_pool"),
|
||||
mutates_args=("key_pool", "value_pool", "stage_key", "stage_value"),
|
||||
)
|
||||
def _paged_write_attn_op(
|
||||
query: torch.Tensor,
|
||||
@@ -535,6 +729,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
|
||||
max_query_len: int,
|
||||
max_seq_len: int,
|
||||
softmax_scale: float,
|
||||
stage_key: torch.Tensor | None,
|
||||
stage_value: torch.Tensor | None,
|
||||
reuse_history: bool,
|
||||
stage_first_block: int,
|
||||
) -> torch.Tensor:
|
||||
return _paged_write_attn_impl(
|
||||
query,
|
||||
@@ -553,6 +751,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
|
||||
max_query_len,
|
||||
max_seq_len,
|
||||
softmax_scale,
|
||||
stage_key,
|
||||
stage_value,
|
||||
reuse_history,
|
||||
stage_first_block,
|
||||
)
|
||||
|
||||
@_paged_write_attn_op.register_fake
|
||||
@@ -573,6 +775,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
|
||||
max_query_len,
|
||||
max_seq_len,
|
||||
softmax_scale,
|
||||
stage_key=None,
|
||||
stage_value=None,
|
||||
reuse_history=False,
|
||||
stage_first_block=0,
|
||||
):
|
||||
return torch.empty_like(query)
|
||||
|
||||
@@ -598,4 +804,8 @@ def paged_write_attn(
|
||||
inputs.max_query_len,
|
||||
inputs.max_seq_len,
|
||||
softmax_scale,
|
||||
inputs.stage_key,
|
||||
inputs.stage_value,
|
||||
inputs.reuse_history,
|
||||
inputs.stage_first_block,
|
||||
)
|
||||
@@ -14,13 +14,14 @@ from vllm.logger import init_logger
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.interface import supports_step_execution
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
|
||||
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput, KVPrefetchJob
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
|
||||
from vllm_omni.experimental.ar_diffusion.capability import (
|
||||
ARDiffusionKVCacheSpec,
|
||||
SupportsARDiffusionPipeline,
|
||||
SupportsARDiffusionWarmup,
|
||||
supports_chunk_step_grouping,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
|
||||
@@ -144,6 +145,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
|
||||
kv_branches=spec.kv_branches,
|
||||
session_capacity=spec.session_capacity,
|
||||
cross_attention_lengths=spec.cross_attention_lengths,
|
||||
cross_attention_kv_heads=spec.cross_attention_kv_heads,
|
||||
frames_per_block=spec.frames_per_block,
|
||||
max_scratch_tokens_per_branch=spec.max_scratch_tokens_per_branch,
|
||||
model_owned_state_bytes_per_session=spec.model_owned_state_bytes_per_session,
|
||||
@@ -163,7 +165,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
|
||||
config.window_chunks,
|
||||
config.sink_chunks,
|
||||
[(kv_branch.name, kv_branch.local_index) for kv_branch in spec.kv_branches],
|
||||
spec.cross_attention_lengths,
|
||||
{
|
||||
name: (length, spec.cross_attention_kv_heads.get(name, spec.num_kv_heads))
|
||||
for name, length in spec.cross_attention_lengths.items()
|
||||
},
|
||||
self._session_capacity,
|
||||
spec.session_capacity,
|
||||
)
|
||||
@@ -338,6 +343,90 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
|
||||
"ARDiffusionModelRunner does not support request-batch execution; use request mode with max_num_seqs=1."
|
||||
)
|
||||
|
||||
def _groups_chunk_steps(self, scheduler_output: DiffusionSchedulerOutput) -> bool:
|
||||
"""Policy: does this stepwise call run the scheduled request's chunk to its boundary?
|
||||
|
||||
Only on the realtime AR path, for a lone streaming request of a
|
||||
pipeline that declares ``supports_chunk_step_grouping``. The session's
|
||||
KV is bound for the whole call and the scheduler regains control at
|
||||
every chunk boundary, which is where this path's interactions are
|
||||
applied anyway; what it gives up is the chance to act between two
|
||||
steps of one chunk.
|
||||
"""
|
||||
return (
|
||||
bool(self.od_config.streaming_output)
|
||||
and len(scheduler_output.scheduled_request_ids) == 1
|
||||
and supports_chunk_step_grouping(self.pipeline)
|
||||
)
|
||||
|
||||
def _execute_stepwise_core(
|
||||
self,
|
||||
scheduler_output: DiffusionSchedulerOutput,
|
||||
*,
|
||||
record_output_peak_memory: bool,
|
||||
in_diffusion_kv_memory_profile: bool = False,
|
||||
) -> BatchRunnerOutput:
|
||||
"""Run the shared single-step core once, or repeatedly until the chunk boundary.
|
||||
|
||||
The shared runner keeps its one-step-per-call semantics. When the
|
||||
grouping policy applies, the remaining steps of the request's current
|
||||
chunk are driven from here with a continuation of the scheduler
|
||||
output: the request is presented as a cached request, with nothing
|
||||
new to admit, nothing to retire and no connector work, so every
|
||||
further step is exactly what the next scheduler cycle would have
|
||||
asked for. Grouping the probes and the commit of one block this way
|
||||
drops the scheduler/executor round trips between them.
|
||||
|
||||
The loop stops at the first call that emits a result (a chunk or an
|
||||
error), finishes the request, or leaves it without state; the base
|
||||
core already handles a mid-chunk failure or interrupt by finishing
|
||||
the request with an error and dropping its state.
|
||||
"""
|
||||
output = super()._execute_stepwise_core(
|
||||
scheduler_output,
|
||||
record_output_peak_memory=record_output_peak_memory,
|
||||
in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
|
||||
)
|
||||
if not self._groups_chunk_steps(scheduler_output):
|
||||
return output
|
||||
request_id = scheduler_output.scheduled_request_ids[0]
|
||||
continuation = dataclasses.replace(
|
||||
scheduler_output,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
|
||||
finished_req_ids=set(),
|
||||
kv_prefetch_job=None,
|
||||
kv_connector_metadata=None,
|
||||
)
|
||||
if not self._chunk_step_pending(output, request_id):
|
||||
return output
|
||||
max_steps = self.state_cache[request_id].chunk_num_steps
|
||||
if max_steps is None:
|
||||
return output
|
||||
steps = 1
|
||||
while self._chunk_step_pending(output, request_id):
|
||||
if steps >= max_steps:
|
||||
self.state_cache.pop(request_id)
|
||||
runner_output = output.get_request_output(request_id)
|
||||
assert runner_output is not None
|
||||
runner_output.finished = True
|
||||
runner_output.result = DiffusionOutput(error=f"Chunk did not complete within {max_steps} denoise steps")
|
||||
break
|
||||
steps += 1
|
||||
output = super()._execute_stepwise_core(
|
||||
continuation,
|
||||
record_output_peak_memory=record_output_peak_memory,
|
||||
in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
|
||||
)
|
||||
return output
|
||||
|
||||
def _chunk_step_pending(self, output: BatchRunnerOutput, request_id: str) -> bool:
|
||||
"""More steps of the request's current chunk remain after ``output``."""
|
||||
runner_output = output.get_request_output(request_id)
|
||||
if runner_output is None or runner_output.finished or runner_output.result is not None:
|
||||
return False
|
||||
return request_id in self.state_cache
|
||||
|
||||
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
|
||||
"""Bind runner-owned KV for one stepwise invocation, then inherit the step loop."""
|
||||
pipeline = getattr(self, "pipeline", None)
|
||||
|
||||
Reference in new issue
Block a user