Skip to content

[Flydsl] Add Qwen-Image VAE classic conv shapes to conv3d tests - #1069

Draft
huizzhan wants to merge 6 commits into
mainfrom
feature/qwenimage-vae-conv-tests
Draft

[Flydsl] Add Qwen-Image VAE classic conv shapes to conv3d tests#1069
huizzhan wants to merge 6 commits into
mainfrom
feature/qwenimage-vae-conv-tests

Conversation

@huizzhan

@huizzhan huizzhan commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

  • Add T2I (T=1) correctness cases for Qwen-Image VAE convs in tests/kernels/test_conv3d_implicit.py (BF16) and tests/kernels/test_conv3d_implicit_fp8.py (FP8).
  • Shapes come from forward-hook traces of AutoencoderKLQwenImage (bf16, all 7 official T2I resolutions), not from assuming the decoder channel ladder halves at every up block.
  • CausalConv3d 3×3×3 is tested via its T=1 2D degeneration (weight[:, :, 2, :, :] + conv2d). Down/upsample paths are native Conv2d.
  • 18 unique shapes × BF16 + 18 × FP8 = 36 cases. There is no torch.float16 suite on this branch.

T=1 degeneration (CausalConv3d 3×3×3)

QwenImageCausalConv3d zeros self.padding and uses causal F.pad:

_padding = (1, 1, 1, 1, 2, 0)   # H/W symmetric; time: pad 2 in front, 0 at the back

For T2I, the logical activation is (N, C, 1, H, W). After causal pad the time axis is [0, 0, x], so only the last kernel slice sees real pixels:

CausalConv3d(x5d, W5d)  ≡  Conv2d(x.squeeze(T), W[:, :, 2, :, :], stride=1, padding=1)

Do not feed the 5D tensors to FlyDSL with symmetric padding=1 — that pads time on both sides and is not causal.

Why decoder ResBlock channels were corrected

QwenImageDecoder3d applies in_dim // 2 inside the UpBlock loop, which cancels the channel halving from upsample2d / upsample3d. There is no 48-channel tensor in this VAE. Decoder ResBlock channel pairs coincide with the encoder ones already on the list; ids name both sides.

Still uncovered on purpose: mid-block attention 1×1 convs (~0.1% of decode MACs) and post_quant_conv. The four (3,1,1) time_conv layers are dead on the T2I path (0 calls across all 7 resolutions).

Test shapes

N=1, T2I T=1. Activation/weight in tests: BF16 or FP8 (e4m3fn); bias is fp32; BF16 reference is F.conv2d.

Freq = number of times this exact (Cin, Cout, H, W) kernel runs in one T2I vae.encode + vae.decode at that resolution (hook-corrected decoder: no 48-ch tensors). 1×1 shortcuts, attention 1×1, quant convs, and dead time_conv are not included. The ten 1024 3×3×3 rows sum to 52 (all CausalConv3d 3×3×3 in the net). *_1328 rows are the same layers at 1328² (H_l = 166), not extra calls on a 1024 run.

1. CausalConv3d 3×3×3 → Conv2d (stride=1, padding=1) — 12

id Cin→Cout Freq Enc / Dec 5D input (N,C,T,H,W) 5D weight 5D output 2D input (tested) 2D weight 2D output
enc_conv_in 3→96 1 1 / 0 (1, 3, 1, 1024, 1024) (96, 3, 3, 3, 3) (1, 96, 1, 1024, 1024) (1, 3, 1024, 1024) (96, 3, 3, 3) (1, 96, 1024, 1024)
enc_e0_res__dec_d3_res 96→96 10 4 / 6 (1, 96, 1, 1024, 1024) (96, 96, 3, 3, 3) (1, 96, 1, 1024, 1024) (1, 96, 1024, 1024) (96, 96, 3, 3) (1, 96, 1024, 1024)
enc_e1_res1 96→192 1 1 / 0 (1, 96, 1, 512, 512) (192, 96, 3, 3, 3) (1, 192, 1, 512, 512) (1, 96, 512, 512) (192, 96, 3, 3) (1, 192, 512, 512)
enc_e1_res2__dec_d2_res 192→192 9 3 / 6 (1, 192, 1, 512, 512) (192, 192, 3, 3, 3) (1, 192, 1, 512, 512) (1, 192, 512, 512) (192, 192, 3, 3) (1, 192, 512, 512)
enc_e2_res1__dec_d1_res1 192→384 2 1 / 1 (1, 192, 1, 256, 256) (384, 192, 3, 3, 3) (1, 384, 1, 256, 256) (1, 192, 256, 256) (384, 192, 3, 3) (1, 384, 256, 256)
enc_e2_res2__dec_d1_res 384→384 8 3 / 5 (1, 384, 1, 256, 256) (384, 384, 3, 3, 3) (1, 384, 1, 256, 256) (1, 384, 256, 256) (384, 384, 3, 3) (1, 384, 256, 256)
enc_e3_mid__dec_mid_d0 384→384 18 8 / 10 (1, 384, 1, 128, 128) (384, 384, 3, 3, 3) (1, 384, 1, 128, 128) (1, 384, 128, 128) (384, 384, 3, 3) (1, 384, 128, 128)
enc_conv_out 384→32 1 1 / 0 (1, 384, 1, 128, 128) (32, 384, 3, 3, 3) (1, 32, 1, 128, 128) (1, 384, 128, 128) (32, 384, 3, 3) (1, 32, 128, 128)
dec_conv_in 16→384 1 0 / 1 (1, 16, 1, 128, 128) (384, 16, 3, 3, 3) (1, 384, 1, 128, 128) (1, 16, 128, 128) (384, 16, 3, 3) (1, 384, 128, 128)
dec_conv_out 96→3 1 0 / 1 (1, 96, 1, 1024, 1024) (3, 96, 3, 3, 3) (1, 3, 1, 1024, 1024) (1, 96, 1024, 1024) (3, 96, 3, 3) (1, 3, 1024, 1024)
dec_bottleneck_1328 384→384 18 8 / 10 (1, 384, 1, 166, 166) (384, 384, 3, 3, 3) (1, 384, 1, 166, 166) (1, 384, 166, 166) (384, 384, 3, 3) (1, 384, 166, 166)
dec_d3_res_hot_1328 96→96 10 4 / 6 (1, 96, 1, 1328, 1328) (96, 96, 3, 3, 3) (1, 96, 1, 1328, 1328) (1, 96, 1328, 1328) (96, 96, 3, 3) (1, 96, 1328, 1328)

Freq breakdown for the busy rows:

  • 96→96 @ full res: E0 Res ×2 (conv1/2 ×4) + D3 Res ×3 (conv1/2 ×6)
  • 192→192 @ H/2: E1 (Res#1 conv2 + Res#2 ×2 = 3) + D2 ×6
  • 384→384 @ H/4: E2 (3) + D1 remaining (5)
  • 384→384 @ H/8: E3 ×4 + enc mid ×4 + dec mid ×4 + D0 ×6 = 18
  • D1 first 3×3×3 is 192→384 (pairs with E2 Res#1 conv1)

2. Downsample Conv2d (ZeroPad2d((0,1,0,1)) + k=3, s=2, p=0) — 3

Native 2D in the VAE; 5D is the feature map before folding T into the batch (N·T = 1 here). Pad is applied outside FlyDSL. Each spatial stage runs once per encode.

id C Freq Enc / Dec 5D feature (N,C,T,H,W) 2D before pad Kernel input (after pad) Weight 2D output 5D output
enc_e0_downsample 96 1 1 / 0 (1, 96, 1, 1024, 1024) (1, 96, 1024, 1024) (1, 96, 1025, 1025) (96, 96, 3, 3) (1, 96, 512, 512) (1, 96, 1, 512, 512)
enc_e1_downsample_spatial 192 1 1 / 0 (1, 192, 1, 512, 512) (1, 192, 512, 512) (1, 192, 513, 513) (192, 192, 3, 3) (1, 192, 256, 256) (1, 192, 1, 256, 256)
enc_e2_downsample_spatial 384 1 1 / 0 (1, 384, 1, 256, 256) (1, 384, 256, 256) (1, 384, 257, 257) (384, 384, 3, 3) (1, 384, 128, 128) (1, 384, 1, 128, 128)

3. Upsample Conv2d (after nearest-exact ×2; k=3, s=1, p=1) — 3

Nearest upsample is outside the kernel. Conv runs at the already-doubled H×W with Conv2d(dim, dim // 2, …). Each stage runs once per decode.

id Cin→Cout Freq Enc / Dec 5D feature (already ×2) 2D input (tested) Weight 2D output 5D output
dec_d0_upsample 384→192 1 0 / 1 (1, 384, 1, 256, 256) (1, 384, 256, 256) (192, 384, 3, 3) (1, 192, 256, 256) (1, 192, 1, 256, 256)
dec_d1_upsample 384→192 1 0 / 1 (1, 384, 1, 512, 512) (1, 384, 512, 512) (192, 384, 3, 3) (1, 192, 512, 512) (1, 192, 1, 512, 512)
dec_d2_upsample 192→96 1 0 / 1 (1, 192, 1, 1024, 1024) (1, 192, 1024, 1024) (96, 192, 3, 3) (1, 96, 1024, 1024) (1, 96, 1, 1024, 1024)

Decode-side MAC share of these upsample convs (from traces): 384→192 @512² ~3.1%, 192→96 @1024² ~3.1%, 384→192 @256² ~0.8%.

Hottest cases by call frequency

Rank id Cin→Cout Spatial 5D input (N,C,T,H,W) Freq Enc / Dec Share of 52 Resolution
1 enc_e3_mid__dec_mid_d0 384→384 128² (1, 384, 1, 128, 128) 18 8 / 10 35% 1024
1 dec_bottleneck_1328 384→384 166² (1, 384, 1, 166, 166) 18 8 / 10 35% 1328 (default)
2 enc_e0_res__dec_d3_res 96→96 1024² (1, 96, 1, 1024, 1024) 10 4 / 6 19% 1024
2 dec_d3_res_hot_1328 96→96 1328² (1, 96, 1, 1328, 1328) 10 4 / 6 19% 1328 (default)
3 enc_e1_res2__dec_d2_res 192→192 512² (1, 192, 1, 512, 512) 9 3 / 6 17% 1024
4 enc_e2_res2__dec_d1_res 384→384 256² (1, 384, 1, 256, 256) 8 3 / 5 15% 1024

Test plan

  • gfx95x: pytest tests/kernels/test_conv3d_implicit.py -k qwenimage
  • gfx95x: pytest tests/kernels/test_conv3d_implicit_fp8.py -k qwenimage
  • Non-CDNA4 hosts skip (@_skip_non_cdna4 / @_skip_no_fp8)
  • BF16: allclose vs F.conv2d (rtol=atol=2e-2)
  • FP8: mean relative error vs bf16 F.conv2d < 2e-2; output dtype bf16

Not in this PR: video T>1 true 5D 3×3×3, feat_cache streaming, (3,1,1) time_conv.

@huizzhan @cursoragent
Add Qwen-Image VAE classic conv shapes to conv3d tests.
e599cc1 
Cover T2I T=1 ResBlock 3x3x3 (via 2D weight slice) and Resample downsample
cases in both BF16 and FP8 implicit-GEMM conv test suites.

Co-authored-by: Cursor <cursoragent@cursor.com>
@huizzhan huizzhan changed the title Add Qwen-Image VAE classic conv shapes to conv3d tests. [Flydsl] Add Qwen-Image VAE classic conv shapes to conv3d tests Aug 26, 2026
huizzhan and others added 3 commits August 26, 2026 11:08
@huizzhan @cursoragent
Correct Qwen-Image VAE decoder conv shapes against hook traces.
eb3d092 
The decoder cases were derived from the config by assuming the channel
count keeps halving across up blocks (384 -> 192 -> 96 -> 48). It does
not: QwenImageDecoder3d applies in_dim // 2 inside the UpBlock loop,
which cancels the halving that upsample2d/3d performs. Forward-hook
traces of AutoencoderKLQwenImage show up_blocks.1 running 192->384 then
384->384, up_blocks.2 running 192->192, and up_blocks.3 running 96->96,
so dec_d1_res, dec_d2_res and dec_d3_res1 were exercising shapes that
never occur in the model. Their real counterparts are already covered by
the encoder entries, which is why dropping them loses no coverage; the
ids now name both sides.

Add the Resample upsample2d/3d convs, which were missing entirely and
are the largest real gap: 384->192 @512 and 192->96 @1024 are ~3% of
decode MACs each.

Co-authored-by: Cursor <cursoragent@cursor.com>
@huizzhan @cursoragent
Stop the conv 1D/2D entry points from repacking weights every launch.
8b1e425 
_prep_weight memoizes the packed weight on id(w), but _conv2d_impl and
_conv1d_impl hand it `weight.reshape(...)` — a fresh view object on every
call — so the cache missed unconditionally and re-ran pad/permute/contiguous
per launch, while _WEIGHT_CACHE accumulated dead weakref entries.

Thread the caller-owned weight down as a cache anchor and prune dead entries
when the cache fills. Measured on gfx950 with the Qwen-Image VAE 384->384
3x3 stride-1 layer: 20 calls now repack once instead of 20 times, 0.178 ->
0.166 ms per call. Every T=1 causal-conv call in that VAE takes the 2D entry
point, so this is ~0.99 ms of a 18.15 ms encode+decode.

Co-authored-by: Cursor <cursoragent@cursor.com>
@huizzhan @cursoragent
Dispatch the conv and transpose launches through the compiled-functio…
bc7c60a 
…n cache.

Both kernels were invoked through their @flyc.jit wrappers, which re-marshal
arguments (DLPack, cache lookup) on every call. That costs ~30 us of host time
per launch irrespective of shape, and with two launches per conv it left the
whole Qwen-Image VAE path host-bound rather than GPU-bound: the 12.6 MB NCHW->
NHWC transpose at the VAE bottleneck measured 38 us while the kernel itself
runs in 9 us.

Route both through _run_compiled so the CompiledFunction is cached after the
first call. Each exe is already memoized per shape by an lru_cache, so a cached
CompiledFunction can only ever be re-invoked on the shape it was compiled for.

Measured on gfx950, per-shape in fresh processes, min of 5 alternating trials.
Weighted by the hook-traced call counts of one 1024 T2I encode+decode
(58 calls): 17.07 -> 13.78 ms, from 1.11x to 1.38x of MIOpen, with 14 of 16
shapes now at or above parity instead of 8. The 384->384 3x3 stride-1 layers
that dominate the network flip from losing to winning: @128 x18 0.62x -> 1.18x,
@256 x8 0.83x -> 1.06x, and @166 x18 (the 1328 default) 0.79x -> 1.21x.

Co-authored-by: Cursor <cursoragent@cursor.com>
@huizzhan
huizzhan force-pushed the feature/qwenimage-vae-conv-tests branch from 255da94 to bc7c60a Compare September 3, 2026 06:34
huizzhan and others added 2 commits September 3, 2026 09:40
@huizzhan @cursoragent
Add Qwen-Image VAE GEMM vs hipBLASLt fair comparison.
124b9e2 
Isolated-process sweep of the 18 T2I conv shapes against MIOpen and
hipBLASLt, plus the script that plots GEMM kernel quality. Default tile
and best-of-five tile are both recorded so the figure can show the gap
on the high-frequency 384→384 layers.

Co-authored-by: Cursor <cursoragent@cursor.com>
@huizzhan @cursoragent
Drop Qwen-Image VAE cases from the FP8 conv3d tests for now.
425f1c1 
Keep the BF16 suite; the FP8 qwenimage parametrization is not needed on this PR yet.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@huizzhan