KV 融合实现完成。改动总结： 速度微弱提升psnr略微上升

attention.py — 3处改动：
  1. __init__ 添加 _kv_fused = False 标志
  2.新增 fuse_kv() 方法：将 to_k + to_v → to_kv，同时处理 _ip/_as/_aa 辅助 KV 对
  2. bmm_forward 两个分支加_kv_fused 判断，用to_kv().chunk(2, dim=-1) 替代分别调用

scripts/evaluation/world_model_interaction.py

@@ -625,6 +625,12 @@ def run_inference(args: argparse.Namespace, gpu_num: int, gpu_no: int) -> None:
     # Compile hot ResBlocks for operator fusion
     apply_torch_compile(model)
 
+    # Fuse KV projections in attention layers (to_k + to_v → to_kv)
+    from unifolm_wma.modules.attention import CrossAttention
+    kv_count = sum(1 for m in model.modules()
+                   if isinstance(m, CrossAttention) and m.fuse_kv())
+    print(f"    ✓ KV fused: {kv_count} attention layers")
+
     # Export precision-converted checkpoint if requested
     if args.export_precision_ckpt:
         export_path = args.export_precision_ckpt

src/unifolm_wma/models/diffusion_head/conditional_unet1d.py

@@ -567,6 +567,11 @@ class ConditionalUnet1D(nn.Module):
         # Broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         timesteps = timesteps.expand(sample.shape[0])
         global_feature = self.diffusion_step_encoder(timesteps)
+        # Pre-expand global_feature once (reused in every down/mid/up block)
+        if self.use_linear_act_proj:
+            global_feature_expanded = global_feature.unsqueeze(1).expand(-1, T, -1)
+        else:
+            global_feature_expanded = global_feature.unsqueeze(1).expand(-1, 2, -1)
         (imagen_cond_down, imagen_cond_mid, imagen_cond_up
          ) = imagen_cond[0:4], imagen_cond[4], imagen_cond[5:]  #NOTE HAND CODE
 
@@ -603,15 +608,11 @@ class ConditionalUnet1D(nn.Module):
 
             if self.use_linear_act_proj:
                 imagen_cond = imagen_cond.reshape(B, T, -1)
-                cur_global_feature = global_feature.unsqueeze(
-                    1).repeat_interleave(repeats=T, dim=1)
             else:
                 imagen_cond = imagen_cond.permute(0, 3, 1, 2)
                 imagen_cond = imagen_cond.reshape(B, 2, -1)
-                cur_global_feature = global_feature.unsqueeze(
-                    1).repeat_interleave(repeats=2, dim=1)
             cur_global_feature = torch.cat(
-                [cur_global_feature, global_cond, imagen_cond], axis=-1)
+                [global_feature_expanded, global_cond, imagen_cond], axis=-1)
             x = resnet(x, cur_global_feature)
             x = resnet2(x, cur_global_feature)
             h.append(x)
@@ -638,15 +639,11 @@ class ConditionalUnet1D(nn.Module):
         imagen_cond = rearrange(imagen_cond, '(b t) c d -> b t c d', b=B)
         if self.use_linear_act_proj:
             imagen_cond = imagen_cond.reshape(B, T, -1)
-            cur_global_feature = global_feature.unsqueeze(1).repeat_interleave(
-                repeats=T, dim=1)
         else:
             imagen_cond = imagen_cond.permute(0, 3, 1, 2)
             imagen_cond = imagen_cond.reshape(B, 2, -1)
-            cur_global_feature = global_feature.unsqueeze(1).repeat_interleave(
-                repeats=2, dim=1)
         cur_global_feature = torch.cat(
-            [cur_global_feature, global_cond, imagen_cond], axis=-1)
+            [global_feature_expanded, global_cond, imagen_cond], axis=-1)
         x = resnet(x, cur_global_feature)
         x = resnet2(x, cur_global_feature)
 
@@ -683,16 +680,12 @@ class ConditionalUnet1D(nn.Module):
 
             if self.use_linear_act_proj:
                 imagen_cond = imagen_cond.reshape(B, T, -1)
-                cur_global_feature = global_feature.unsqueeze(
-                    1).repeat_interleave(repeats=T, dim=1)
             else:
                 imagen_cond = imagen_cond.permute(0, 3, 1, 2)
                 imagen_cond = imagen_cond.reshape(B, 2, -1)
-                cur_global_feature = global_feature.unsqueeze(
-                    1).repeat_interleave(repeats=2, dim=1)
 
             cur_global_feature = torch.cat(
-                [cur_global_feature, global_cond, imagen_cond], axis=-1)
+                [global_feature_expanded, global_cond, imagen_cond], axis=-1)
 
             x = torch.cat((x, h.pop()), dim=1)
             x = resnet(x, cur_global_feature)

src/unifolm_wma/models/samplers/ddim.py

@@ -251,6 +251,13 @@ class DDIMSampler(object):
         dp_ddim_scheduler_action.set_timesteps(len(timesteps))
         dp_ddim_scheduler_state.set_timesteps(len(timesteps))
         ts = torch.empty((b, ), device=device, dtype=torch.long)
+        noise_buf = torch.empty_like(img)
+        # Pre-convert schedule arrays to inference dtype (avoid per-step .to())
+        _dtype = img.dtype
+        _alphas = (self.model.alphas_cumprod if ddim_use_original_steps else self.ddim_alphas).to(_dtype)
+        _alphas_prev = (self.model.alphas_cumprod_prev if ddim_use_original_steps else self.ddim_alphas_prev).to(_dtype)
+        _sqrt_one_minus = (self.model.sqrt_one_minus_alphas_cumprod if ddim_use_original_steps else self.ddim_sqrt_one_minus_alphas).to(_dtype)
+        _sigmas = (self.ddim_sigmas_for_original_num_steps if ddim_use_original_steps else self.ddim_sigmas).to(_dtype)
         enable_cross_attn_kv_cache(self.model)
         enable_ctx_cache(self.model)
         try:
@@ -286,6 +293,8 @@ class DDIMSampler(object):
                     x0=x0,
                     fs=fs,
                     guidance_rescale=guidance_rescale,
+                    noise_buf=noise_buf,
+                    schedule_arrays=(_alphas, _alphas_prev, _sqrt_one_minus, _sigmas),
                     **kwargs)
 
                 img, pred_x0, model_output_action, model_output_state = outs
@@ -339,6 +348,8 @@ class DDIMSampler(object):
                       mask=None,
                       x0=None,
                       guidance_rescale=0.0,
+                      noise_buf=None,
+                      schedule_arrays=None,
                       **kwargs):
         b, *_, device = *x.shape, x.device
 
@@ -384,16 +395,18 @@ class DDIMSampler(object):
             e_t = score_corrector.modify_score(self.model, e_t, x, t, c,
                                                **corrector_kwargs)
 
-        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        if schedule_arrays is not None:
-        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+            alphas, alphas_prev, sqrt_one_minus_alphas, sigmas = schedule_arrays
-        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        else:
-        sigmas = self.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+            alphas = (self.model.alphas_cumprod if use_original_steps else self.ddim_alphas).to(x.dtype)
+            alphas_prev = (self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev).to(x.dtype)
+            sqrt_one_minus_alphas = (self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas).to(x.dtype)
+            sigmas = (self.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas).to(x.dtype)
 
-        # Use 0-d tensors directly (already on device); broadcasting handles shape
+        a_t = alphas[index]
-        a_t = alphas[index].to(x.dtype)
+        a_prev = alphas_prev[index]
-        a_prev = alphas_prev[index].to(x.dtype)
+        sigma_t = sigmas[index]
-        sigma_t = sigmas[index].to(x.dtype)
+        sqrt_one_minus_at = sqrt_one_minus_alphas[index]
-        sqrt_one_minus_at = sqrt_one_minus_alphas[index].to(x.dtype)
 
         if self.model.parameterization != "v":
             pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
@@ -411,8 +424,12 @@ class DDIMSampler(object):
 
         dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
 
-        noise = sigma_t * noise_like(x.shape, device,
+        if noise_buf is not None:
-                                     repeat_noise) * temperature
+            noise_buf.normal_()
+            noise = sigma_t * noise_buf * temperature
+        else:
+            noise = sigma_t * noise_like(x.shape, device,
+                                         repeat_noise) * temperature
         if noise_dropout > 0.:
             noise = torch.nn.functional.dropout(noise, p=noise_dropout)
 

src/unifolm_wma/modules/attention.py

@@ -99,6 +99,7 @@ class CrossAttention(nn.Module):
         self.agent_action_context_len = agent_action_context_len
         self._kv_cache = {}
         self._kv_cache_enabled = False
+        self._kv_fused = False
 
         self.cross_attention_scale_learnable = cross_attention_scale_learnable
         if self.image_cross_attention:
@@ -116,6 +117,27 @@ class CrossAttention(nn.Module):
                 self.register_parameter('alpha_caa',
                                         nn.Parameter(torch.tensor(0.)))
 
+    def fuse_kv(self):
+        """Fuse to_k/to_v into to_kv (2 Linear → 1). Works for all layers."""
+        k_w = self.to_k.weight  # (inner_dim, context_dim)
+        v_w = self.to_v.weight
+        self.to_kv = nn.Linear(k_w.shape[1], k_w.shape[0] * 2, bias=False)
+        self.to_kv.weight = nn.Parameter(torch.cat([k_w, v_w], dim=0))
+        del self.to_k, self.to_v
+        if self.image_cross_attention:
+            for suffix in ('_ip', '_as', '_aa'):
+                k_attr = f'to_k{suffix}'
+                v_attr = f'to_v{suffix}'
+                kw = getattr(self, k_attr).weight
+                vw = getattr(self, v_attr).weight
+                fused = nn.Linear(kw.shape[1], kw.shape[0] * 2, bias=False)
+                fused.weight = nn.Parameter(torch.cat([kw, vw], dim=0))
+                setattr(self, f'to_kv{suffix}', fused)
+                delattr(self, k_attr)
+                delattr(self, v_attr)
+        self._kv_fused = True
+        return True
+
     def forward(self, x, context=None, mask=None):
         spatial_self_attn = (context is None)
         k_ip, v_ip, out_ip = None, None, None
@@ -276,14 +298,20 @@ class CrossAttention(nn.Module):
                                     self.agent_action_context_len +
                                     self.text_context_len:, :]
 
-            k = self.to_k(context_ins)
+            if self._kv_fused:
-            v = self.to_v(context_ins)
+                k, v = self.to_kv(context_ins).chunk(2, dim=-1)
-            k_ip = self.to_k_ip(context_image)
+                k_ip, v_ip = self.to_kv_ip(context_image).chunk(2, dim=-1)
-            v_ip = self.to_v_ip(context_image)
+                k_as, v_as = self.to_kv_as(context_agent_state).chunk(2, dim=-1)
-            k_as = self.to_k_as(context_agent_state)
+                k_aa, v_aa = self.to_kv_aa(context_agent_action).chunk(2, dim=-1)
-            v_as = self.to_v_as(context_agent_state)
+            else:
-            k_aa = self.to_k_aa(context_agent_action)
+                k = self.to_k(context_ins)
-            v_aa = self.to_v_aa(context_agent_action)
+                v = self.to_v(context_ins)
+                k_ip = self.to_k_ip(context_image)
+                v_ip = self.to_v_ip(context_image)
+                k_as = self.to_k_as(context_agent_state)
+                v_as = self.to_v_as(context_agent_state)
+                k_aa = self.to_k_aa(context_agent_action)
+                v_aa = self.to_v_aa(context_agent_action)
 
             q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h),
                           (q, k, v))
@@ -304,8 +332,11 @@ class CrossAttention(nn.Module):
         else:
             if not spatial_self_attn:
                 context = context[:, :self.text_context_len, :]
-            k = self.to_k(context)
+            if self._kv_fused:
-            v = self.to_v(context)
+                k, v = self.to_kv(context).chunk(2, dim=-1)
+            else:
+                k = self.to_k(context)
+                v = self.to_v(context)
 
             q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h),
                           (q, k, v))

src/unifolm_wma/modules/networks/wma_model.py

@@ -690,6 +690,8 @@ class WMAModel(nn.Module):
         self._ctx_cache = {}
         # fs_embed cache
         self._fs_embed_cache = None
+        # Pre-created CUDA stream for parallel action/state UNet
+        self._side_stream = torch.cuda.Stream() if not self.base_model_gen_only else None
 
     def forward(self,
                 x: Tensor,
@@ -849,8 +851,8 @@ class WMAModel(nn.Module):
         if not self.base_model_gen_only:
             ba, _, _ = x_action.shape
             ts_state = timesteps[:ba] if b > 1 else timesteps
-            # Run action_unet and state_unet in parallel via CUDA streams
+            # Run action_unet and state_unet in parallel via pre-created CUDA stream
-            s_stream = torch.cuda.Stream()
+            s_stream = self._side_stream
             s_stream.wait_stream(torch.cuda.current_stream())
             with torch.cuda.stream(s_stream):
                 s_y = self.state_unet(x_state, ts_state, hs_a,

unitree_z1_dual_arm_cleanup_pencils/case1/output.log

@@ -1,14 +1,14 @@
 /mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
   __import__("pkg_resources").declare_namespace(__name__)
-2026-02-10 10:36:44.797852: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-02-10 17:57:48.047156: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
-2026-02-10 10:36:44.801300: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-10 17:57:48.050303: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
-2026-02-10 10:36:44.837891: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+2026-02-10 17:57:48.081710: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
-2026-02-10 10:36:44.837946: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+2026-02-10 17:57:48.081741: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
-2026-02-10 10:36:44.839880: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+2026-02-10 17:57:48.083577: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
-2026-02-10 10:36:44.849073: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
+2026-02-10 17:57:48.091772: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
-2026-02-10 10:36:44.849365: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+2026-02-10 17:57:48.092045: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
 To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
-2026-02-10 10:36:45.644793: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
+2026-02-10 17:57:48.787960: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
 [rank: 0] Global seed set to 123
 /mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
   @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
@@ -41,6 +41,7 @@ INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k).
     ⚠ Found 601 fp32 params, converting to bf16
     ✓ All parameters converted to bfloat16
     ✓ torch.compile: 3 ResBlocks in output_blocks[5, 8, 9]
+    ✓ KV fused: 66 attention layers
 INFO:root:***** Configing Data *****
 >>> unitree_z1_stackbox: 1 data samples loaded.
 >>> unitree_z1_stackbox: data stats loaded.
@@ -116,7 +117,7 @@ DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing WmfImagePlugin
 DEBUG:PIL.Image:Importing XbmImagePlugin
 DEBUG:PIL.Image:Importing XpmImagePlugin
-DEBUG:PIL.Image:Importing XVThumbImagePlugin
+DEBUG:PIL.Image:Importing XVThumbImagePlugin
 
  12%|█▎        | 1/8 [01:03<07:22, 63.25s/it]
  25%|██▌       | 2/8 [02:02<06:05, 60.93s/it]
@@ -140,6 +141,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin
 >>> Step 4: generating actions ...
 >>> Step 4: interacting with world model ...
 >>>>>>>>>>>>>>>>>>>>>>>>
->>> Step 5: generating actions ...
+>>> Step 5: generating actions ...
->>> Step 5: interacting with world model ...
+>>> Step 5: interacting with world model ...
->>>>>>>>>>>>>>>>>>>>>>>>
+>>>>>>>>>>>>>>>>>>>>>>>>

unitree_z1_dual_arm_cleanup_pencils/case1/psnr_result1.json

@@ -1,5 +1,5 @@
 {
     "gt_video": "/mnt/ASC1637/unifolm-world-model-action/unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/unitree_z1_dual_arm_cleanup_pencils_case1_amd.mp4",
     "pred_video": "/mnt/ASC1637/unifolm-world-model-action/unitree_z1_dual_arm_cleanup_pencils/case1/output/inference/0_full_fs4.mp4",
-    "psnr": 31.802224855380352
+    "psnr": 32.442113263955434
 }