[Feature]: conv3d_implicit — accept and emit NDHWC so chained convs skip the layout round-trip

[ZJLi2013]

Suggestion Description

Ask

conv3d_implicit takes NCDHW and calls _ncdhw_to_ndhwc on every invocation (L705). In a network that is a chain of 3D convolutions, every intermediate could have stayed in NDHWC. Please add a layout parameter so a caller can convert twice per network instead of twice per layer:

conv3d_implicit(
    x, weight, bias=None, stride=1, padding=0,
    layout="NCDHW",        # "NCDHW" (default, unchanged) | "NDHWC"
    out_layout=None,       # None -> same as `layout`
    **kwargs,
)

layout="NDHWC" skips the input transpose; out_layout="NDHWC" lets the non-split-K epilogue write (n, do, ho, wo, k), which is a change of output view rather than new kernel work.

Why

On a Wan 3D causal VAE encoder (bf16, input (1, 3, 33, 352, 640), 278 conv calls, MI350X / gfx950 / ROCm 7.1.1), layout movement is 17.4% of the 193 ms encoder — 992 batched_transpose kernels, 33.6 ms. Swapping MIOpen for conv3d_implicit already removes the weight and output round-trips (_prep_weight caches, the splitk==1 epilogue writes NCDHW), and is 1.20x faster overall on these shapes: 108.2 → 89.9 ms measured per-shape. What remains is the per-call input transpose, roughly 20 ms that nothing in the network needs.

The dominant shapes are rank-5 in=[1,96,6,354,642] w=[96,96,3,3,3], in=[1,384,3,46,82] w=[384,384,3,3,3] and rank-4 in=[4,96,353,641] w=[96,96,3,3], all bf16. Happy to test a branch on them.

Related

• [Issue]: Tune conv perf #861 [Issue]: Tune conv perf — tuning the conv itself; this request is about the data movement around it.

Operating System

No response

GPU

No response

ROCm Component

No response

[jhinpan]

A second consumer for this: MiniMax H3 (Hailuo 3.0)

MiniMaxAI/MiniMax-H3 (open weights, Aug 2026) ships a visual VAE whose encoder is exactly this pattern — a causal 3D CNN consuming NCDHW. A layout parameter would apply there too, but the shapes are far enough from the Wan ones above that tuning for Wan alone won't cover it. Details below in case they help scope the change.

Paths below are Ref2VA/video_vae/...; the repo ships the VAE once per checkpoint and FL2VA/video_vae/ is an identical copy.

The call site is already a clean drop-in

Ref2VA/video_vae/conv.py does its own padding and calls F.conv3d with padding=0:

class BaseConv3d(nn.Conv3d):
    def forward(self, x):
        if sum(self.padding) == 0:
            return super().forward(x)
        x = self._apply_padding(x)   # spatial reflect pad + temporal causal pad
        return F.conv3d(x, self.weight, self.bias,
                        stride=self.stride, padding=0, dilation=self.dilation)

Every conv in the encoder is constructed as SpatialParallelConv3d, which inherits this forward unchanged and only overrides _apply_padding. Spatial padding is reflect, temporal is causal single-sided zero pad, both applied before the conv — so nothing here depends on conv3d_implicit's own padding handling. (The 1x1x1 convs take the sum(self.padding) == 0 branch into nn.Conv3d.forward instead, but they are negligible FLOPs.)

Shapes

From Ref2VA/video_vae/source/config.json: ch=128, ch_mult=[1,2,2,4,4,8], num_res_blocks=2, space_down=[2,2,2,2,1,1], time_down=[1,2,2,1,1,1], use_3d_conv=true, causal_encoder=true. Tiling is configured one level up, in Ref2VA/video_vae/config.json: vae_encoder_tiling=1, vae_tile_size=256, vae_tile_overlap_min=64, vae_clip_length=17. stack_tiling is absent from both files and defaults to False, so every tile runs at N=1.

30 3x3x3 conv3d calls per tile per 17-frame clip, bf16:

stage	shape	count
conv_in	in=[1,3,17,256,256] w=[128,3,3,3,3]	1
L0	in=[1,128,17,256,256] w=[128,128,3,3,3]	4, + 1 at stride (1,2,2)
L1	in=[1,128,17,128,128] w=[256,128,3,3,3] then in=[1,256,17,128,128] w=[256,256,3,3,3]	1 + 3, + 1 at stride (2,2,2)
L2	in=[1,256,9,64,64] w=[256,256,3,3,3]	4, + 1 at stride (2,2,2)
L3	in=[1,256,5,32,32] w=[512,256,3,3,3] then in=[1,512,5,32,32] w=[512,512,3,3,3]	1 + 3, + 1 at stride (1,2,2)
L4	in=[1,512,5,16,16] w=[512,512,3,3,3]	4
L5	in=[1,512,5,16,16] w=[1024,512,3,3,3] then in=[1,1024,5,16,16] w=[1024,1024,3,3,3]	1 + 3
conv_out	in=[1,1024,5,16,16] w=[48,1024,3,3,3]	1

Plus three 1x1x1 nin_shortcut convs and the 1x1x1 quant_conv.

FLOPs concentrate in L0 (~1 TFLOP per conv). L4/L5 are the opposite regime — 16x16 spatial with 512/1024 channels. Neither overlaps much with the w=[96,96,3,3,3] / w=[384,384,3,3,3] shapes being tuned in #861.

Call volume is higher than the Wan profile because of the 256-tile decomposition. Note that split_tiles starts at ceil(input_len / tile_size) and then increments until tiles minus overlaps actually cover the input, so 768 lands on 4 tiles and 1376 on 7, not 3 and 6. A 5 s 768x1376 reference clip is 120 frames at 24 fps, snapped up to the next 17n + 5 = 124 — so 4x7 = 28 spatial tiles and 8 temporal chunks (7 full plus one of 5). Roughly 30 x 28 x 8 = 6.7k conv3d calls at N=1, against 278 for the Wan encoder above. Same per-call input transpose, ~24x the count.

An extra argument for NDHWC that isn't in the original ask

This encoder is not a pure chain of convs — every pair is separated by a GroupNorm, and the config sets use_t_isolated_gn=true, so Ref2VA/video_vae/norm.py runs a time-to-batch reshape on the way in:

def _merge_time_to_batch(x):        # N,C,D,H,W -> N*D,C,1,H,W
    batch, channels, depth, height, width = x.shape
    return x.permute(0, 2, 1, 3, 4).contiguous().view(batch * depth, channels, 1, height, width)

and permutes back on the way out, where _split_time_from_batch is .contiguous() as well. With 12 ResnetBlock3D x 2 norms plus norm_out = 25 norms, that is 50 full-tensor copies per encoder forward, on top of the conv layout transposes.

Worth stressing: these copies survive the fusion H3 already ships. norm.py has an opt-in MINIMAX_H3_USE_FUSED_NORM=true path that swaps in FusedGroupNorm3D to fuse the GroupNorm with its SiLU — and that path still calls _merge_time_to_batch on the way in and _split_time_from_batch on the way out. So these 50 copies are not something kernel fusion can reach on this network; only a layout change can.

This cuts both ways:

• Against the ask as filed. Adding layout to conv3d_implicit alone does not let intermediates stay in NDHWC here. The norm wants C at dim 1, so the transpose just relocates from the conv into the norm.
• For a stronger version of the ask. In NDHWC, _merge_time_to_batch becomes a pure view: D is already adjacent to N, so (N,D,H,W,C) -> (N*D,H,W,C) is free, and all 50 copies disappear. What that needs is a channels-last GroupNorm — and in NDHWC a group's C/32 channels are contiguous at each spatial position, so the reduction is memory-friendly rather than strided.

So on this network the end-to-end win requires conv and norm to move together. It may be worth saying that explicitly in the API contract: out_layout="NDHWC" pays off only when the caller has channels-last versions of whatever sits between the convs.

Scope caveats

Three things bound how much of H3 this touches:

• The decoder is not a CNN. H3 uses ViT3DDecoder (36 layers), so the decode path gets nothing from this. Unlike Wan, conv3d exists only on the encode side.
• The encoder only runs when there is visual input to encode. Pure t2va never calls it. fl2va calls it on 1–2 frames (D=1). ref2va — up to 9 reference images and 3 reference video clips — is the path that produces the profile above.
• Spatial parallelism interacts with the layout choice. The profile above is single-GPU (vae_encoder_parallel=0). With it enabled, SpatialParallelConv3d._apply_padding exchanges halos across ranks along vae_chunk_dim=-1, i.e. W. In NDHWC, W is no longer the innermost dimension, so the exchanged borders stop being contiguous. Worth deciding up front whether layout is meant to compose with spatial parallelism or is single-rank only.

Happy to profile the H3 encoder on MI350X the same way if a branch lands.

[jiacao-amd]

Thanks both — this is well-scoped and the profiles help a lot.

I'm on board with adding the layout parameter. A few notes from the kernel side on what's cheap vs. what needs real work:

Input side (layout="NDHWC") — trivial, it's just skipping the _ncdhw_to_ndhwc call. No kernel change.

Output side (out_layout="NDHWC") — splits in two:

• split-K path: already free. The kernel materializes (npq, k), which is NDHWC-flat, then does .view(n,do,ho,wo,k).permute(0,4,1,2,3) at the end. So out_layout="NDHWC" there is deleting a permute. @ZJLi2013 is right that it's a view change, not new work.
• splitk==1 path: needs a real epilogue change. It allocates (n,k,do,ho,wo) and scatters into it, so NDHWC means recomputing the store offset. Modest work, and it should end up faster — the store becomes contiguous along k instead of strided.

On the GroupNorm point (@jhinpan) — good catch, agreed. Conv-only won't get you there on H3, since the transpose just moves into the norm. If GroupNorm goes channels-last too, this gets much better.

On spatial parallelism — hadn't thought about that, thanks. You're right that NDHWC breaks the contiguous halo along W. I'd lean toward single-rank only for a first cut.