From 7fe6d899a4c5c6c146ca1d4d627b4a73410b6382 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Shiqiao=20Gu=20=28=E8=B0=B7=E7=9F=B3=E6=A1=A5=29?= <77222802+gushiqiao@users.noreply.github.com> Date: Tue, 14 Apr 2026 13:03:41 +0800 Subject: [PATCH] Adjust kv_cache_size calculation based on local_attn_size --- .../networks/wan/infer/lingbot_fast/transformer_infer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lightx2v/models/networks/wan/infer/lingbot_fast/transformer_infer.py b/lightx2v/models/networks/wan/infer/lingbot_fast/transformer_infer.py index f5111186..2d33679d 100755 --- a/lightx2v/models/networks/wan/infer/lingbot_fast/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/lingbot_fast/transformer_infer.py @@ -52,7 +52,11 @@ def _initialize_kv_cache(self, dtype, device): return kv_cache1 = [] ws = self._sp_world_size() - self.kv_cache_size = self._kv_size // ws + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length // ws + else: + kv_cache_size = self._kv_size // ws + self.kv_cache_size = kv_cache_size n, d = self.num_heads, self.head_dim if self.kv_quant_config is not None: