[IvenHsu01]

Adds ck.temporal_interleave_upsample_nearest2x(x, interleave=False), which takes a [B, C, T, H, W] tensor, moves the frames into the batch ('b c t h w -> (b t) c h w') and upsamples it 2x with nearest-exact, returning the contiguous [B * T, C, 2H, 2W] input of the Conv2d that follows. With interleave=True, x is the [B, 2C, T, H, W] output of the temporal conv and the frames of its two channel halves are interleaved first ([B, C, 2T, H, W], frame 2t + k from half k). This is the upsample body of the Wan VAE Resample (comfy/ldm/wan/vae.py and comfy/ldm/wan/vae2_2.py, Resample.forward, modes upsample3d and upsample2d): the reshape / torch.stack / reshape frame interleave after time_conv, the rearrange to (b t) c h w and nn.Upsample(scale_factor=(2., 2.), mode='nearest-exact') before the Conv2d in self.resample. In ComfyUI's current path the interleave form runs as four kernels that each read and write the full activation (three copy kernels for the stack and the rearranged view, then the upsample), and the plain form with more than one frame copies the rearranged view before the upsample; comfy-kitchen has no upsample or interleave op.

• eager: the reference sequence (reshape / stack / reshape, permute + reshape for the rearrange, F.interpolate(..., scale_factor=(2., 2.), mode="nearest-exact")), made contiguous.
• triton: one gather kernel writes the contiguous output straight from x. A program row loads a tile of source rows through x's strides, duplicates every pixel along the row with tl.interleave and stores it to both output rows. For a scale of 2, nearest-exact maps an output index d to d // 2, so the output is a copy of input pixels, bit-identical to eager. fp32 / fp16 / bf16 x on cuda devices (CUDA and ROCm); any strides. The tile shape is autotuned over 7 (block_r, block_ws) shapes with 4 / 8 warps. The call rule declines empty inputs and, for the plain form (interleave=False), outputs below 12M elements, where the single eager upsample kernel was faster than a direct launch of the kernel (crossover between 9.58M and 14.4M output elements on the RTX 5090 and the Radeon 8060S); those calls run eager.

Speedup over ComfyUI's current path (reshape / stack / reshape + rearrange + nn.Upsample nearest-exact), bf16, per call of the public op at the input shapes captured from real Wan VAE decodes. Each arm runs in a fresh process, ABAB with 2 pairs (20 warmup calls, then 5 loops of 100 back-to-back calls with a synchronize around each loop, median per call), so the CPU launch and dispatch cost of each call counts. Times are pair 1; the speedups are pair 1, pair 2. Every output is torch.equal to the current path, and every call was served by the backend the call rule picks (eager for the first row, triton for the rest):

shape (calls per decode)	output elements	RTX 5090, current path -> kitchen	Radeon 8060S, current path -> kitchen
Wan2.1 upsample3d plain, x (1, 384, 1, 60, 104) (1, served eager)	9.58M	14.6 -> 23.4 us (0.62x, 0.70x)	129.2 -> 149.4 us (0.87x, 0.85x)
Wan2.1 upsample3d plain, x (1, 384, 1, 120, 208) (1)	38.34M	49.3 -> 32.7 us (1.51x, 1.50x)	705.2 -> 519.0 us (1.36x, 1.39x)
Wan2.1 upsample3d interleave, x (1, 768, 2, 60, 104) (4)	38.34M	100.4 -> 52.8 us (1.90x, 1.90x)	1113.6 -> 550.1 us (2.02x, 2.02x)
Wan2.1 upsample3d interleave, x (1, 768, 2, 120, 208) (8)	153.35M	575.0 -> 273.8 us (2.10x, 2.11x)	4184.8 -> 1869.9 us (2.24x, 2.24x)
Wan2.1 upsample2d, x (1, 192, 1, 240, 416) (1)	76.68M	192.8 -> 138.4 us (1.39x, 1.39x)	1320.4 -> 927.4 us (1.42x, 1.43x)
Wan2.1 upsample2d, x (1, 192, 2, 240, 416) (16)	153.35M	473.1 -> 274.1 us (1.73x, 1.73x)	3369.7 -> 1930.7 us (1.75x, 1.74x)
Wan2.2 upsample3d plain, x (1, 1024, 1, 44, 80) (1)	14.42M	23.2 -> 34.4 us (0.67x, 0.68x)	300.4 -> 232.2 us (1.29x, 1.36x)
Wan2.2 upsample3d plain, x (1, 1024, 1, 88, 160) (1)	57.67M	159.9 -> 101.4 us (1.58x, 1.58x)	1039.2 -> 711.6 us (1.46x, 1.48x)
Wan2.2 upsample3d interleave, x (1, 2048, 1, 44, 80) (30)	28.84M	72.0 -> 49.6 us (1.45x, 2.18x)	900.7 -> 415.7 us (2.17x, 2.18x)
Wan2.2 upsample3d interleave, x (1, 2048, 2, 88, 160) (30)	230.69M	934.5 -> 409.9 us (2.28x, 2.28x)	6298.6 -> 2887.8 us (2.18x, 2.21x)
Wan2.2 upsample2d, x (1, 512, 1, 176, 320) (1)	115.34M	312.8 -> 203.0 us (1.54x, 1.54x)	1926.9 -> 1323.6 us (1.46x, 1.45x)
Wan2.2 upsample2d, x (1, 512, 4, 176, 320) (30)	461.37M	1542.4 -> 826.6 us (1.87x, 1.87x)	9808.2 -> 5477.1 us (1.79x, 1.79x)

On the RTX 5090 the kitchen time of the (1, 2048, 1, 44, 80) interleave shape was 49.6 us in pair 1 and 33.0 us in pair 2. Two plain-form shapes are slower than the current path through the public op. The (1, 384, 1, 60, 104) shape runs eager by design, but the eager fallback pays the op's dispatch cost too (on both GPUs). The (1, 1024, 1, 44, 80) shape (14.42M output elements) is above the 12M threshold, so triton serves it: 1.29x / 1.36x on the Radeon 8060S, but 0.67x / 0.68x on the RTX 5090. The threshold is shared by both GPUs and is not changed in this PR.

Environment: RTX 5090 on Windows (torch 2.14.1+cu130, CUDA 13.0, triton-windows 3.9.0, Python 3.12); Radeon 8060S (RDNA3.5) on Windows (ROCm torch 2.15.0a0+rocm10.1.0a20260907, triton-windows 3.8.0, Python 3.12). Tests and benchmarks on 4f24e06 on both GPUs (54673ea and the head 668b547 after it only edit comments and the text of one call-rule message): test_temporal_interleave_upsample_nearest2x.py 64, test_backends.py 12, test_constraints.py 42 passed on the RTX 5090 with the CUDA extension from the 0.2.41 wheel; the same on the Radeon 8060S with the HIP extension only, except test_constraints.py 40 passed / 2 skipped (CUDA-only).

Not covered yet: end to end in ComfyUI (not run on either GPU); Linux; fp16 / fp32 timing (the benchmark is bf16 only). Interleave inputs smaller than the Wan decode shapes (smallest measured output 28.8M elements) are served by triton without a size limit but were not timed; the 12M routing threshold was measured in bf16 only. The single-image path of the Wan2.2 VAE (feat_cache None) runs the upsample and the Conv2d together in row strips (strip_apply) and is not a call site of this op. ComfyUI integration (Wan VAE) is a separate PR; ComfyUI keeps comfy-kitchen's triton backend off unless the server is started with --enable-triton-backend.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts