From e599cc1cb467a995b72ce2a7efe222041533f39c Mon Sep 17 00:00:00 2001 From: huizzhan Date: Wed, 26 Aug 2026 10:00:04 +0000 Subject: [PATCH 1/2] Add Qwen-Image VAE classic conv shapes to conv3d tests. 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 --- tests/kernels/test_conv3d_implicit.py | 65 ++++++++++++++++++ tests/kernels/test_conv3d_implicit_fp8.py | 80 +++++++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/tests/kernels/test_conv3d_implicit.py b/tests/kernels/test_conv3d_implicit.py index a2b5dcd69..59daeea8b 100644 --- a/tests/kernels/test_conv3d_implicit.py +++ b/tests/kernels/test_conv3d_implicit.py @@ -254,3 +254,68 @@ def test_conv1d_vs_torch(s, stride, padding): assert y.shape == y_ref.shape assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +# ---- Qwen-Image VAE classic conv (T2I T=1) --------------------------------- +# CausalConv3d 3x3x3 degenerates to conv2d with weight[:, :, 2, :, :]. + + +_QWENIMAGE_T1_RES3 = [ + pytest.param(3, 96, 1024, 1024, id="enc_conv_in"), + pytest.param(96, 96, 1024, 1024, id="enc_e0_res"), + pytest.param(96, 192, 512, 512, id="enc_e1_res1"), + pytest.param(192, 192, 512, 512, id="enc_e1_res2"), + pytest.param(192, 384, 256, 256, id="enc_e2_res1"), + pytest.param(384, 384, 256, 256, id="enc_e2_res2"), + pytest.param(384, 384, 128, 128, id="enc_e3_mid_bottleneck"), + pytest.param(384, 32, 128, 128, id="enc_conv_out"), + pytest.param(16, 384, 128, 128, id="dec_conv_in"), + pytest.param(384, 384, 128, 128, id="dec_mid_d0"), + pytest.param(192, 192, 256, 256, id="dec_d1_res"), + pytest.param(96, 96, 512, 512, id="dec_d2_res"), + pytest.param(48, 96, 1024, 1024, id="dec_d3_res1"), + pytest.param(96, 96, 1024, 1024, id="dec_d3_res_hot"), + pytest.param(96, 3, 1024, 1024, id="dec_conv_out"), + pytest.param(384, 384, 166, 166, id="dec_bottleneck_1328"), + pytest.param(96, 96, 1328, 1328, id="dec_d3_res_hot_1328"), +] + +_QWENIMAGE_DOWN2D = [ + pytest.param(96, 1024, 1024, id="enc_e0_downsample"), + pytest.param(192, 512, 512, id="enc_e1_downsample_spatial"), + pytest.param(384, 256, 256, id="enc_e2_downsample_spatial"), +] + + +@_skip_non_cdna4 +@pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_T1_RES3) +def test_qwenimage_vae_t1_res3_bf16(c_in, c_out, h, w): + torch.manual_seed(8800 + c_in + c_out + h + w) + x2 = torch.randn((1, c_in, h, w), device="cuda", dtype=torch.bfloat16) + weight5 = torch.randn((c_out, c_in, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + weight2 = weight5[:, :, 2, :, :] + bias = torch.randn((c_out,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x2, weight2, bias=bias, stride=1, padding=1) + y_ref = F.conv2d(x2, weight2, bias=bias.to(torch.bfloat16), stride=1, padding=1) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape == (1, c_out, h, w) + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +@_skip_non_cdna4 +@pytest.mark.parametrize("c,h,w", _QWENIMAGE_DOWN2D) +def test_qwenimage_vae_downsample2d_bf16(c, h, w): + torch.manual_seed(9100 + c + h + w) + x2 = torch.randn((1, c, h, w), device="cuda", dtype=torch.bfloat16) + x_pad = F.pad(x2, (0, 1, 0, 1)) + weight = torch.randn((c, c, 3, 3), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((c,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x_pad, weight, bias=bias, stride=2, padding=0) + y_ref = F.conv2d(x_pad, weight, bias=bias.to(torch.bfloat16), stride=2, padding=0) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) diff --git a/tests/kernels/test_conv3d_implicit_fp8.py b/tests/kernels/test_conv3d_implicit_fp8.py index a43ba2864..7d1327432 100644 --- a/tests/kernels/test_conv3d_implicit_fp8.py +++ b/tests/kernels/test_conv3d_implicit_fp8.py @@ -202,3 +202,83 @@ def test_fp8_transpose_dword_aligned_spatial(c, t, h, w): ref = x.permute(0, 2, 3, 4, 1).contiguous().view(torch.int8).view(-1) assert torch.equal(got, ref) + + +# ---- Qwen-Image VAE classic conv (T2I T=1) --------------------------------- +# Same 2D-degenerate shapes as test_conv3d_implicit.test_qwenimage_vae_*. + + +_QWENIMAGE_T1_RES3 = [ + pytest.param(3, 96, 1024, 1024, id="enc_conv_in"), + pytest.param(96, 96, 1024, 1024, id="enc_e0_res"), + pytest.param(96, 192, 512, 512, id="enc_e1_res1"), + pytest.param(192, 192, 512, 512, id="enc_e1_res2"), + pytest.param(192, 384, 256, 256, id="enc_e2_res1"), + pytest.param(384, 384, 256, 256, id="enc_e2_res2"), + pytest.param(384, 384, 128, 128, id="enc_e3_mid_bottleneck"), + pytest.param(384, 32, 128, 128, id="enc_conv_out"), + pytest.param(16, 384, 128, 128, id="dec_conv_in"), + pytest.param(384, 384, 128, 128, id="dec_mid_d0"), + pytest.param(192, 192, 256, 256, id="dec_d1_res"), + pytest.param(96, 96, 512, 512, id="dec_d2_res"), + pytest.param(48, 96, 1024, 1024, id="dec_d3_res1"), + pytest.param(96, 96, 1024, 1024, id="dec_d3_res_hot"), + pytest.param(96, 3, 1024, 1024, id="dec_conv_out"), + pytest.param(384, 384, 166, 166, id="dec_bottleneck_1328"), + pytest.param(96, 96, 1328, 1328, id="dec_d3_res_hot_1328"), +] + +_QWENIMAGE_DOWN2D = [ + pytest.param(96, 1024, 1024, id="enc_e0_downsample"), + pytest.param(192, 512, 512, id="enc_e1_downsample_spatial"), + pytest.param(384, 256, 256, id="enc_e2_downsample_spatial"), +] + + +@_skip_no_fp8 +@pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_T1_RES3) +def test_qwenimage_vae_t1_res3_fp8(c_in, c_out, h, w): + torch.manual_seed(8900 + c_in + c_out + h + w) + x2 = _fp8(torch.randn((1, c_in, h, w), device="cuda", dtype=torch.bfloat16)) + weight5 = _fp8(torch.randn((c_out, c_in, 3, 3, 3), device="cuda", dtype=torch.bfloat16)) + weight2 = weight5[:, :, 2, :, :] + bias = torch.randn((c_out,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit_fp8(x2, weight2, bias=bias, stride=1, padding=1) + y_ref = F.conv2d( + x2.to(torch.bfloat16), + weight2.to(torch.bfloat16), + bias=bias.to(torch.bfloat16), + stride=1, + padding=1, + ) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert y.dtype == torch.bfloat16 + rel = (y.float() - y_ref.float()).abs().mean() / y_ref.float().abs().mean().clamp_min(1e-6) + assert rel.item() < 2e-2, f"Qwen-Image FP8 rel_err {rel.item():.3e}" + + +@_skip_no_fp8 +@pytest.mark.parametrize("c,h,w", _QWENIMAGE_DOWN2D) +def test_qwenimage_vae_downsample2d_fp8(c, h, w): + torch.manual_seed(9200 + c + h + w) + x2 = _fp8(torch.randn((1, c, h, w), device="cuda", dtype=torch.bfloat16)) + x_pad = F.pad(x2, (0, 1, 0, 1)) + weight = _fp8(torch.randn((c, c, 3, 3), device="cuda", dtype=torch.bfloat16)) + bias = torch.randn((c,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit_fp8(x_pad, weight, bias=bias, stride=2, padding=0) + y_ref = F.conv2d( + x_pad.to(torch.bfloat16), + weight.to(torch.bfloat16), + bias=bias.to(torch.bfloat16), + stride=2, + padding=0, + ) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + rel = (y.float() - y_ref.float()).abs().mean() / y_ref.float().abs().mean().clamp_min(1e-6) + assert rel.item() < 2e-2, f"Qwen-Image FP8 downsample rel_err {rel.item():.3e}" From eb3d092906b8fc2da7e33bbaec0de8e6f3c2d9cc Mon Sep 17 00:00:00 2001 From: huizzhan Date: Wed, 26 Aug 2026 11:08:52 +0000 Subject: [PATCH 2/2] Correct Qwen-Image VAE decoder conv shapes against hook traces. 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 --- tests/kernels/test_conv3d_implicit.py | 47 ++++++++++++++++++----- tests/kernels/test_conv3d_implicit_fp8.py | 45 +++++++++++++++++----- 2 files changed, 72 insertions(+), 20 deletions(-) diff --git a/tests/kernels/test_conv3d_implicit.py b/tests/kernels/test_conv3d_implicit.py index 59daeea8b..cd3595152 100644 --- a/tests/kernels/test_conv3d_implicit.py +++ b/tests/kernels/test_conv3d_implicit.py @@ -258,34 +258,45 @@ def test_conv1d_vs_torch(s, stride, padding): # ---- Qwen-Image VAE classic conv (T2I T=1) --------------------------------- # CausalConv3d 3x3x3 degenerates to conv2d with weight[:, :, 2, :, :]. +# Shapes below are the 1024x1024 spatial ladder plus the two hottest layers of +# the 1328x1328 default resolution, taken from forward-hook traces of +# AutoencoderKLQwenImage rather than from the config alone: the decoder halves +# its channel count inside the UpBlock loop (in_dim // 2) before each stage, so +# its ResBlock channel pairs coincide with the encoder ones instead of +# continuing 384 -> 192 -> 96 -> 48. _QWENIMAGE_T1_RES3 = [ pytest.param(3, 96, 1024, 1024, id="enc_conv_in"), - pytest.param(96, 96, 1024, 1024, id="enc_e0_res"), + pytest.param(96, 96, 1024, 1024, id="enc_e0_res__dec_d3_res"), pytest.param(96, 192, 512, 512, id="enc_e1_res1"), - pytest.param(192, 192, 512, 512, id="enc_e1_res2"), - pytest.param(192, 384, 256, 256, id="enc_e2_res1"), - pytest.param(384, 384, 256, 256, id="enc_e2_res2"), - pytest.param(384, 384, 128, 128, id="enc_e3_mid_bottleneck"), + pytest.param(192, 192, 512, 512, id="enc_e1_res2__dec_d2_res"), + pytest.param(192, 384, 256, 256, id="enc_e2_res1__dec_d1_res1"), + pytest.param(384, 384, 256, 256, id="enc_e2_res2__dec_d1_res"), + pytest.param(384, 384, 128, 128, id="enc_e3_mid__dec_mid_d0"), pytest.param(384, 32, 128, 128, id="enc_conv_out"), pytest.param(16, 384, 128, 128, id="dec_conv_in"), - pytest.param(384, 384, 128, 128, id="dec_mid_d0"), - pytest.param(192, 192, 256, 256, id="dec_d1_res"), - pytest.param(96, 96, 512, 512, id="dec_d2_res"), - pytest.param(48, 96, 1024, 1024, id="dec_d3_res1"), - pytest.param(96, 96, 1024, 1024, id="dec_d3_res_hot"), pytest.param(96, 3, 1024, 1024, id="dec_conv_out"), pytest.param(384, 384, 166, 166, id="dec_bottleneck_1328"), pytest.param(96, 96, 1328, 1328, id="dec_d3_res_hot_1328"), ] +# Resample downsample2d: ZeroPad2d((0, 1, 0, 1)) then Conv2d(k=3, s=2, p=0). _QWENIMAGE_DOWN2D = [ pytest.param(96, 1024, 1024, id="enc_e0_downsample"), pytest.param(192, 512, 512, id="enc_e1_downsample_spatial"), pytest.param(384, 256, 256, id="enc_e2_downsample_spatial"), ] +# Resample upsample2d/3d: nearest-exact x2 happens outside the kernel, so the +# conv runs at the already-doubled resolution with Conv2d(dim, dim // 2, k=3, +# s=1, p=1). +_QWENIMAGE_UP2D = [ + pytest.param(384, 192, 256, 256, id="dec_d0_upsample"), + pytest.param(384, 192, 512, 512, id="dec_d1_upsample"), + pytest.param(192, 96, 1024, 1024, id="dec_d2_upsample"), +] + @_skip_non_cdna4 @pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_T1_RES3) @@ -319,3 +330,19 @@ def test_qwenimage_vae_downsample2d_bf16(c, h, w): assert y.shape == y_ref.shape assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +@_skip_non_cdna4 +@pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_UP2D) +def test_qwenimage_vae_upsample2d_bf16(c_in, c_out, h, w): + torch.manual_seed(9400 + c_in + c_out + h + w) + x2 = torch.randn((1, c_in, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((c_out, c_in, 3, 3), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((c_out,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x2, weight, bias=bias, stride=1, padding=1) + y_ref = F.conv2d(x2, weight, bias=bias.to(torch.bfloat16), stride=1, padding=1) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape == (1, c_out, h, w) + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) diff --git a/tests/kernels/test_conv3d_implicit_fp8.py b/tests/kernels/test_conv3d_implicit_fp8.py index 7d1327432..a9891fdb2 100644 --- a/tests/kernels/test_conv3d_implicit_fp8.py +++ b/tests/kernels/test_conv3d_implicit_fp8.py @@ -210,19 +210,14 @@ def test_fp8_transpose_dword_aligned_spatial(c, t, h, w): _QWENIMAGE_T1_RES3 = [ pytest.param(3, 96, 1024, 1024, id="enc_conv_in"), - pytest.param(96, 96, 1024, 1024, id="enc_e0_res"), + pytest.param(96, 96, 1024, 1024, id="enc_e0_res__dec_d3_res"), pytest.param(96, 192, 512, 512, id="enc_e1_res1"), - pytest.param(192, 192, 512, 512, id="enc_e1_res2"), - pytest.param(192, 384, 256, 256, id="enc_e2_res1"), - pytest.param(384, 384, 256, 256, id="enc_e2_res2"), - pytest.param(384, 384, 128, 128, id="enc_e3_mid_bottleneck"), + pytest.param(192, 192, 512, 512, id="enc_e1_res2__dec_d2_res"), + pytest.param(192, 384, 256, 256, id="enc_e2_res1__dec_d1_res1"), + pytest.param(384, 384, 256, 256, id="enc_e2_res2__dec_d1_res"), + pytest.param(384, 384, 128, 128, id="enc_e3_mid__dec_mid_d0"), pytest.param(384, 32, 128, 128, id="enc_conv_out"), pytest.param(16, 384, 128, 128, id="dec_conv_in"), - pytest.param(384, 384, 128, 128, id="dec_mid_d0"), - pytest.param(192, 192, 256, 256, id="dec_d1_res"), - pytest.param(96, 96, 512, 512, id="dec_d2_res"), - pytest.param(48, 96, 1024, 1024, id="dec_d3_res1"), - pytest.param(96, 96, 1024, 1024, id="dec_d3_res_hot"), pytest.param(96, 3, 1024, 1024, id="dec_conv_out"), pytest.param(384, 384, 166, 166, id="dec_bottleneck_1328"), pytest.param(96, 96, 1328, 1328, id="dec_d3_res_hot_1328"), @@ -234,6 +229,12 @@ def test_fp8_transpose_dword_aligned_spatial(c, t, h, w): pytest.param(384, 256, 256, id="enc_e2_downsample_spatial"), ] +_QWENIMAGE_UP2D = [ + pytest.param(384, 192, 256, 256, id="dec_d0_upsample"), + pytest.param(384, 192, 512, 512, id="dec_d1_upsample"), + pytest.param(192, 96, 1024, 1024, id="dec_d2_upsample"), +] + @_skip_no_fp8 @pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_T1_RES3) @@ -282,3 +283,27 @@ def test_qwenimage_vae_downsample2d_fp8(c, h, w): assert y.shape == y_ref.shape rel = (y.float() - y_ref.float()).abs().mean() / y_ref.float().abs().mean().clamp_min(1e-6) assert rel.item() < 2e-2, f"Qwen-Image FP8 downsample rel_err {rel.item():.3e}" + + +@_skip_no_fp8 +@pytest.mark.parametrize("c_in,c_out,h,w", _QWENIMAGE_UP2D) +def test_qwenimage_vae_upsample2d_fp8(c_in, c_out, h, w): + torch.manual_seed(9500 + c_in + c_out + h + w) + x2 = _fp8(torch.randn((1, c_in, h, w), device="cuda", dtype=torch.bfloat16)) + weight = _fp8(torch.randn((c_out, c_in, 3, 3), device="cuda", dtype=torch.bfloat16)) + bias = torch.randn((c_out,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit_fp8(x2, weight, bias=bias, stride=1, padding=1) + y_ref = F.conv2d( + x2.to(torch.bfloat16), + weight.to(torch.bfloat16), + bias=bias.to(torch.bfloat16), + stride=1, + padding=1, + ) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert y.dtype == torch.bfloat16 + rel = (y.float() - y_ref.float()).abs().mean() / y_ref.float().abs().mean().clamp_min(1e-6) + assert rel.item() < 2e-2, f"Qwen-Image FP8 upsample rel_err {rel.item():.3e}"