From d6a3551a6200d9053a79b4b4869a9644bfecec87 Mon Sep 17 00:00:00 2001 From: Ding Zuhao Date: Mon, 23 Feb 2026 23:58:22 +0800 Subject: [PATCH] [Bug] [Bagel] Fix kv transfer bug (#1437) Signed-off-by: Ding Zuhao Co-authored-by: Wang Zhipeng: princepride --- vllm_omni/diffusion/models/bagel/bagel_transformer.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 90a64720c..541a03716 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1286,8 +1286,14 @@ class Bagel(nn.Module): @staticmethod def _merge_naive_caches(caches: list) -> NaiveCache: """Merge multiple NaiveCache objects by concatenating KV tensors per layer.""" - merged = NaiveCache(caches[0].num_layers) - for layer_idx in range(merged.num_layers): + if not caches: + # Handle empty list case gracefully if desired, + # though original code also crashed on this. + return NaiveCache(0) + + num_layers = len(caches[0].key_cache) + merged = NaiveCache(num_layers) + for layer_idx in range(num_layers): merged.key_cache[layer_idx] = torch.cat([c.key_cache[layer_idx] for c in caches], dim=0) merged.value_cache[layer_idx] = torch.cat([c.value_cache[layer_idx] for c in caches], dim=0) return merged