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

[Bugfix]Fix Hunyuan Image3 denoise flow alignment (#3857)

Browse Source 
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
bjf-frzandHongsheng Liu authored and GitHub committed 2026-05-27 19:18:01 +08:00
1 parent c349f62839
commit c33ef34654
5 files changed
+246 -133
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
+131 -71
View File
@@ -23,6 +23,7 @@ NUM_HEADS = 4
NUM_KV_HEADS = 2
HEAD_DIM = 16
IMAGE_TOKEN_LEN = 8
SUFFIX_TOKEN_LEN = 3
SCALING = 1.0 / math.sqrt(HEAD_DIM)
@@ -79,6 +80,10 @@ def _make_known_kv(num_tokens, base=0.0):
return k, v
def _gen_timestep_index(bs: int, current_start: int) -> torch.Tensor:
return torch.full((bs, 1), current_start, dtype=torch.long)
def _call_mgr(
mgr,
bs,
@@ -90,6 +95,8 @@ def _call_mgr(
uncond_cfg_prefill=False,
num_image_tokens=IMAGE_TOKEN_LEN,
shard_image_size=None,
gen_timestep_scatter_index=None,
position_ids=None,
):
query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -104,6 +111,8 @@ def _call_mgr(
uncond_cfg_prefill=uncond_cfg_prefill,
num_image_tokens=num_image_tokens,
shard_image_size=shard_image_size,
gen_timestep_scatter_index=gen_timestep_scatter_index,
position_ids=position_ids,
)
@@ -117,54 +126,65 @@ def test_no_ar_kv(bs):
"""
No AR KV injected. Tests the basic first_step cache → update reuse path.
Sequence layout per batch on first_step (q_len=12, IMAGE_TOKEN_LEN=8):
[prompt(3) | image(8) | eoi(1)]
image_size = IMAGE_TOKEN_LEN + 1 = 9
cached_prompt_len = seq_len - image_size = 12 - 9 = 3
Sequence layout per batch on first_step (q_len=14, IMAGE_TOKEN_LEN=8):
[prompt(3) | current timestep/image(8) | suffix(3)]
gen_timestep_scatter_index = 3
cached_prompt_len = 3
After first_step:
image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
image_kv_cache_map stores prompt tokens [0:3] for each batch.
Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + eoi(1) + image(8) = 12):
_reuse_prompt_kv produces [cached_prompt(3) | new_image(8) | zero_eoi(1)] per batch.
Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + current(8) = 11):
_reuse_prompt_kv produces [cached_prompt(3) | new_current(8)] per batch.
"""
mgr = _make_cache_mgr()
assert mgr.image_kv_cache_map is None
# --- first_step ---
q_len = 12
prompt_len = 3
q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
_call_mgr(mgr, bs, q_len=q_len, seq_len=q_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
_call_mgr(
mgr,
bs,
q_len=q_len,
seq_len=q_len,
key_flat=k_flat,
value_flat=v_flat,
first_step=True,
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
)
cached_key, cached_value = mgr.image_kv_cache_map
# 3 prompt tokens cached per batch
assert cached_key.shape[0] == 3 * bs
assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
assert torch.equal(mgr.image_kv_cache_lens, torch.full((bs,), prompt_len))
for b in range(bs):
flat_offset = b * q_len
cache_offset = b * 3
assert torch.allclose(cached_key[cache_offset : cache_offset + 3], k_flat[flat_offset : flat_offset + 3])
assert torch.allclose(cached_value[cache_offset : cache_offset + 3], v_flat[flat_offset : flat_offset + 3])
assert torch.allclose(cached_key[b, :prompt_len], k_flat[flat_offset : flat_offset + prompt_len])
assert torch.allclose(cached_value[b, :prompt_len], v_flat[flat_offset : flat_offset + prompt_len])
# --- update step ---
img_q_len = IMAGE_TOKEN_LEN
update_seq_len = 3 + 1 + img_q_len # cached_prompt(3) + eoi(1) + image(8) = 12
update_seq_len = prompt_len + img_q_len # cached_prompt(3) + current(8) = 11
new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
position_ids = torch.arange(prompt_len, prompt_len + img_q_len).repeat(bs, 1)
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs, position_ids=position_ids)
assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
for b in range(bs):
cache_offset = b * 3
img_offset = b * img_q_len
# Cached prompt preserved
assert torch.allclose(result_k[b, :3], cached_key[cache_offset : cache_offset + 3])
assert torch.allclose(result_k[b, :prompt_len], cached_key[b, :prompt_len])
# New image tokens
assert torch.allclose(result_k[b, 3 : 3 + img_q_len], new_img_k[img_offset : img_offset + img_q_len])
# Zero eoi at the end
assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
assert torch.allclose(
result_k[b, prompt_len : prompt_len + img_q_len],
new_img_k[img_offset : img_offset + img_q_len],
)
# ============================================================
@@ -178,16 +198,17 @@ def test_ar_kv_no_cfg(sp_size):
AR KV injected, bs=1, no CFG. Tests AR prefix prepend + cache + reuse.
sp_size=1:
first_step input: q_len=12, ar_len=5, seq_len=17
Sequence layout: [ar(5) | prompt(3) | image(8) | eoi(1)]
image_size = IMAGE_TOKEN_LEN + 1 = 9
cached_prompt_len = seq_len - image_size = 17 - 9 = 8 (= ar(5) + prompt(3))
first_step input: q_len=14, ar_len=5, seq_len=19
Sequence layout: [ar(5) | prompt(3) | current(8) | suffix(3)]
gen_timestep_scatter_index = 3
cached_prompt_len = ar(5) + prompt(3) = 8
Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
Result: [cached(8) | new_image(8) | zero_eoi(1)]
Update step: q_len=8, seq_len = cached_prompt(8) + current(8) = 16
Result: [cached(8) | new_current(8)]
sp_size=2:
shard_image_size = 4 (passed externally, simulating SP sharding)
gen_timestep_scatter_index = 8
cached_prompt_len = seq_len - shard_image_size = 17 - 4 = 13 (= ar(5) + prompt(8))
Note: with SP, more of the sequence is treated as "prompt" because
image tokens are sharded across ranks.
@@ -204,15 +225,23 @@ def test_ar_kv_no_cfg(sp_size):
mgr._injected_ar_kv = [(ar_k.clone(), ar_v.clone())]
# --- first_step: call _cache_prompt_kv directly ---
bs, q_len = 1, 12
seq_len = q_len + ar_len # 17
bs = 1
q_len = 12 if sp_size > 1 else 14
seq_len = q_len + ar_len
k_raw, v_raw = _make_known_kv(q_len, base=1.0)
key_4d = k_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
val_4d = v_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
shard_image_size = 4 if sp_size > 1 else None
current_start = 8 if sp_size > 1 else 3
mgr.image_kv_cache_map = None
mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
mgr._cache_prompt_kv(
key_4d,
val_4d,
seq_len,
shard_image_size,
gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
)
cached_key, cached_value = mgr.image_kv_cache_map
if sp_size == 1:
@@ -222,25 +251,32 @@ def test_ar_kv_no_cfg(sp_size):
# cached = seq_len - shard_image_size = 17 - 4 = 13
expected_cached_len = seq_len - shard_image_size
assert cached_key.shape[0] == expected_cached_len
assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
# AR KV always at the front
assert torch.allclose(cached_key[:ar_len], ar_k)
assert torch.allclose(cached_value[:ar_len], ar_v)
assert torch.allclose(cached_key[0, :ar_len], ar_k)
assert torch.allclose(cached_value[0, :ar_len], ar_v)
# Prompt tokens follow AR
prompt_cached = expected_cached_len - ar_len
assert torch.allclose(cached_key[ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
assert torch.allclose(cached_key[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
# AR KV consumed
assert mgr._injected_ar_kv is None
# --- update step: call _reuse_prompt_kv directly ---
if sp_size == 1:
img_q_len = IMAGE_TOKEN_LEN
update_seq_len = expected_cached_len + 1 + img_q_len # 8 + 1 + 8 = 17
update_seq_len = expected_cached_len + img_q_len # 8 + 8 = 16
new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=1)
position_ids = torch.arange(expected_cached_len, expected_cached_len + img_q_len).reshape(1, img_q_len)
result_k, result_v = mgr._reuse_prompt_kv(
key_input,
val_input,
update_seq_len,
bs=1,
position_ids=position_ids,
)
assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
# AR + prompt preserved
@@ -248,8 +284,6 @@ def test_ar_kv_no_cfg(sp_size):
assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
# New image tokens
assert torch.allclose(result_k[0, expected_cached_len : expected_cached_len + img_q_len], new_img_k)
# Zero eoi
assert torch.allclose(result_k[0, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
else:
# SP path: _reuse_prompt_kv returns only cached prompt (no image concat)
img_q_len = shard_image_size # 4
@@ -292,13 +326,12 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
→ _injected_ar_kv becomes [(pos_ar(10), pos_av(10)), (neg_k(10), neg_v(10))]
first_step (bs=2, q_len=12, seq_len=22):
Batch 0 (pos): [pos_ar(10) | prompt(3) | image(8) | eoi(1)]
Batch 1 (neg): [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
cached_prompt_len per batch = 22 - 9 = 13
Total cached = 13 * 2 = 26
Batch 0 (pos): [pos_ar(10) | prompt(3) | current(8) | suffix(3)]
Batch 1 (neg): [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
cached_prompt_len per batch = 10 + 3 = 13
Update (bs=2, seq_len = 13 + 1 + 8 = 22):
Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
Update (bs=2, seq_len = 13 + 8 = 21):
Result per batch: [cached(13) | new_current(8)]
CFG Parallel (cfg_parallel=True, bs=1):
This rank handles only the negative branch.
@@ -308,11 +341,11 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
→ _injected_ar_kv = [(neg_k(10), neg_v(10))]
first_step (bs=1, q_len=12, seq_len=22):
[neg_ar(10) | prompt(3) | image(8) | eoi(1)]
cached_prompt_len = 22 - 9 = 13
[neg_ar(10) | prompt(3) | current(8) | suffix(3)]
cached_prompt_len = 13
Update (bs=1, seq_len = 13 + 1 + 8 = 22):
Result: [cached(13) | new_image(8) | zero_eoi(1)]
Update (bs=1, seq_len = 13 + 8 = 21):
Result: [cached(13) | new_current(8)]
"""
positive_reuse_len = 10
negative_reuse_len = 6
@@ -351,46 +384,62 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
# --- first_step ---
q_len = 12
prompt_len = 3
q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
seq_len = q_len + positive_reuse_len # 22
k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
_call_mgr(mgr, bs, q_len=q_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
_call_mgr(
mgr,
bs,
q_len=q_len,
seq_len=seq_len,
key_flat=k_flat,
value_flat=v_flat,
first_step=True,
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
)
cached_key, cached_value = mgr.image_kv_cache_map
cached_prompt_len_per_batch = seq_len - (IMAGE_TOKEN_LEN + 1) # 22 - 9 = 13
assert cached_key.shape[0] == cached_prompt_len_per_batch * bs
cached_prompt_len_per_batch = positive_reuse_len + prompt_len
assert cached_key.shape == (bs, cached_prompt_len_per_batch, NUM_KV_HEADS, HEAD_DIM)
assert mgr._injected_ar_kv is None
if not cfg_parallel:
# bs=2: batch 0 = pos, batch 1 = neg
# Batch 0: pos_ar(10) + prompt(3)
assert torch.allclose(cached_key[:positive_reuse_len], pos_ar_k)
assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
assert torch.allclose(cached_key[0, :positive_reuse_len], pos_ar_k)
assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
# Batch 1: neg_ar(10) + prompt(3)
assert torch.allclose(cached_key[13:23], neg_ar_k)
assert torch.allclose(cached_key[23:26], k_flat[q_len : q_len + 3])
assert torch.allclose(cached_key[1, :positive_reuse_len], neg_ar_k)
assert torch.allclose(cached_key[1, positive_reuse_len:13], k_flat[q_len : q_len + 3])
else:
# bs=1: only neg branch
assert torch.allclose(cached_key[:positive_reuse_len], neg_ar_k)
assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
assert torch.allclose(cached_key[0, :positive_reuse_len], neg_ar_k)
assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
# --- update step ---
img_q_len = IMAGE_TOKEN_LEN
update_seq_len = cached_prompt_len_per_batch + 1 + img_q_len # 13 + 1 + 8 = 22
update_seq_len = cached_prompt_len_per_batch + img_q_len # 13 + 8 = 21
new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
position_ids = torch.arange(cached_prompt_len_per_batch, cached_prompt_len_per_batch + img_q_len).repeat(bs, 1)
result_k, result_v = mgr._reuse_prompt_kv(
key_input,
val_input,
update_seq_len,
bs=bs,
position_ids=position_ids,
)
assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
for b in range(bs):
cache_offset = b * cached_prompt_len_per_batch
# Cached prompt preserved
assert torch.allclose(
result_k[b, :cached_prompt_len_per_batch],
cached_key[cache_offset : cache_offset + cached_prompt_len_per_batch],
cached_key[b, :cached_prompt_len_per_batch],
)
# New image tokens
img_offset = b * img_q_len
@@ -398,8 +447,6 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
result_k[b, cached_prompt_len_per_batch : cached_prompt_len_per_batch + img_q_len],
new_img_k[img_offset : img_offset + img_q_len],
)
# Zero eoi
assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
# ============================================================
@@ -412,7 +459,7 @@ def test_cross_request_isolation():
Verify leftover image_kv_cache_map from a previous request is NOT treated as AR KV.
Setup: mgr has stale cache from a prior request (9 tokens, base=999).
New request: first_step with q_len=12, no AR KV.
New request: first_step with q_len=14, no AR KV.
Expected: stale cache is overwritten. New cache = prompt tokens from current request.
The stale values (999.x) must NOT appear in the new cache.
@@ -421,20 +468,33 @@ def test_cross_request_isolation():
# Simulate leftover from previous request
leftover_k, leftover_v = _make_known_kv(9, base=999.0)
mgr.image_kv_cache_map = (leftover_k, leftover_v)
mgr.image_kv_cache_map = (
leftover_k.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
leftover_v.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
)
assert mgr._injected_ar_kv is None
# New request first_step
bs, seq_len = 1, 12
bs = 1
prompt_len = 3
seq_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
k_flat, v_flat = _make_known_kv(bs * seq_len, base=1.0)
_call_mgr(mgr, bs, q_len=seq_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
_call_mgr(
mgr,
bs,
q_len=seq_len,
seq_len=seq_len,
key_flat=k_flat,
value_flat=v_flat,
first_step=True,
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
)
cached_key, cached_value = mgr.image_kv_cache_map
# cached_prompt_len = 12 - 9 = 3
assert cached_key.shape[0] == 3
assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
# Must be from current request, not stale
assert torch.allclose(cached_key[:3], k_flat[:3])
assert torch.allclose(cached_value[:3], v_flat[:3])
assert torch.allclose(cached_key[0, :prompt_len], k_flat[:prompt_len])
assert torch.allclose(cached_value[0, :prompt_len], v_flat[:prompt_len])
# Stale values must not be present
assert not torch.any(cached_key >= 999.0)