Restore front-zero temporal pad for multi-frame causal conv; two-step F.pad (reflect spatial + constant T)

src/h3_blackwell_runtime/vae_encoder.py

@@ -58,13 +58,19 @@ def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding):
     if x.shape[2] == 1:
         # Keyframe path (matches reference's `autopad="causal_zero"`): truncate
         # the temporal taps to the center 1 and run as an effective 2D conv.
-        # We emulate this with F.conv3d by slicing the kernel's T dim to 1 and
-        # running conv3d with stride_t=1 and no T padding -- output T stays 1.
         half = (kernel_size - 1) // 2
-        kernel_5d = weight[:, :, half:half + 1, :, :]  # (O, I, 1, k, k)
+        kernel_5d = weight[:, :, half:half + 1, :, :]
         return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (spatial_padding, spatial_padding, spatial_padding))
+    if spatial_padding > 0:
+        return _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding)
+    return F.conv3d(x, weight, bias, stride, (0, 0, 0))
+
+
+def _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding):
+    """Two-step pad (reference CausalConv3d): spatial reflect, temporal front-zero."""
+    x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
     x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0))
-    return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))
+    return F.conv3d(x, weight, bias, stride, (0, 0, 0))
 
 
 def _group_norm_3d(x, weight, bias):