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

[Diffusion] Add paged KV cache support for HunyuanImage3 DiT on NPU and GPU (#6563)

Browse Source 
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <2831474076@qq.com>
This commit is contained in:
Wang fuyinzwhzzz0821zwhzzz0821
 authored and GitHub committed 2026-08-31 10:15:46 +00:00
1 parent 6f16f9546d
commit fe4a2a7477
28 files changed
+1789 -205
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py
+11 -2
View File
@@ -154,8 +154,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
torch.manual_seed(1)
query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
key = torch.randn_like(query)
value = torch.randn_like(query)
packed_qkv = torch.randn(
19,
3 * _NUM_HEADS * _HEAD_SIZE,
dtype=torch.float16,
device=device,
)
_, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
assert not key.is_contiguous()
assert not value.is_contiguous()
omni_backend = FlashAttentionImpl(
num_heads=_NUM_HEADS,
head_size=_HEAD_SIZE,