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

kernels/conv/conv3d_implicit.py

@@ -20,6 +20,7 @@
 from flydsl.expr.rocdl.universal import make_buffer_ptr
 from flydsl.expr.typing import T
 from kernels.common.mem_ops import buffer_atomic_add
+from kernels.common.tensor_shim import _run_compiled
 
 TILE_K = 32
 STAGES = 2
@@ -43,21 +44,32 @@ def _autotune_enabled():
 
 
 _WEIGHT_CACHE = {}
+_WEIGHT_CACHE_MAX = 64
 
 
 def _pad_channels(c):
     return (c + LDG_VEC - 1) // LDG_VEC * LDG_VEC
 
 
-def _prep_weight(w, k, kt, kh, kw, c):
-    key = id(w)
+def _prep_weight(w, k, kt, kh, kw, c, anchor=None):
+    """Pack (K,C,T,R,S) -> (K, T*R*S*Cpad), memoized on the caller's tensor.
+
+    ``anchor`` is the weight the caller owns. The 1D/2D entry points reshape to
+    5D before calling, and that view is a fresh object per call, so keying on it
+    would miss every time and repack the weights on every launch.
+    """
+    src = w if anchor is None else anchor
+    key = id(src)
     ent = _WEIGHT_CACHE.get(key)
-    if ent is not None and ent[0]() is w:
+    if ent is not None and ent[0]() is src:
         return ent[1]
     cp = _pad_channels(c)
     wsrc = torch.nn.functional.pad(w, (0, 0, 0, 0, 0, 0, 0, cp - c)) if cp != c else w
     wk = wsrc.permute(0, 2, 3, 4, 1).contiguous().reshape(k, kt * kh * kw * cp)
-    _WEIGHT_CACHE[key] = (weakref.ref(w), wk)
+    if len(_WEIGHT_CACHE) >= _WEIGHT_CACHE_MAX:
+        for dead in [k2 for k2, e in _WEIGHT_CACHE.items() if e[0]() is None]:
+            del _WEIGHT_CACHE[dead]
+    _WEIGHT_CACHE[key] = (weakref.ref(src), wk)
     return wk
 
 
@@ -185,7 +197,10 @@ def _ncdhw_to_ndhwc(x, stream):
         return x.permute(0, 2, 3, 4, 1).contiguous()
     out = torch.empty((n, t, h, w, c), device=x.device, dtype=x.dtype)
     exe = compile_transpose_ncdhw_ndhwc(n, c, s)
-    exe(out, x, torch.cuda.current_stream() if stream is None else stream)
+    # Fast dispatch: the @flyc.jit wrapper re-marshals arguments on every call,
+    # which costs ~30 us of host time and dominates this kernel outright — the
+    # 12.6 MB transpose in the VAE bottleneck runs in 9 us.
+    _run_compiled(exe, out, x, torch.cuda.current_stream() if stream is None else stream)
     return out
 
 
@@ -716,7 +731,9 @@ def _resolve_splitk(splitk, npq, crs, k, device, tile=DEFAULT_TILE):
     return sk
 
 
-def _conv3d_impl(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None, tile=None, autotune=None):
+def _conv3d_impl(
+    x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None, tile=None, autotune=None, weight_src=None
+):
     n, c, d, h, w = x.shape
     k, wc, kt, kh, kw = weight.shape
     assert c == wc
@@ -752,7 +769,7 @@ def _conv3d_impl(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=
     bias_arg = bias.to(torch.float32).contiguous() if has_bias else torch.empty(1, device=x.device, dtype=torch.float32)
 
     x_ndhwc = _ncdhw_to_ndhwc(x, stream)
-    w_packed = _prep_weight(weight, k, kt, kh, kw, wc)
+    w_packed = _prep_weight(weight, k, kt, kh, kw, wc, anchor=weight_src)
 
     shape = (n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias)
 
@@ -765,7 +782,7 @@ def _run(the_tile, the_wgm=1):
         exe = compile_conv3d_implicit(
             n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias, sk, the_tile, the_wgm
         )
-        exe(y, x_ndhwc, w_packed, bias_arg, launch_stream)
+        _run_compiled(exe, y, x_ndhwc, w_packed, bias_arg, launch_stream)
         return y, sk
 
     if tile is not None:
@@ -798,7 +815,7 @@ def _conv2d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs):
     k, wc, r, s = weight.shape
     x5 = x.reshape(n, c, 1, h, w)
     w5 = weight.reshape(k, wc, 1, r, s)
-    y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, sh, sw), padding=(0, ph, pw), **kwargs)
+    y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, sh, sw), padding=(0, ph, pw), weight_src=weight, **kwargs)
     return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[3], y5.shape[4])
 
 
@@ -810,7 +827,7 @@ def _conv1d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs):
     k, wc, s = weight.shape
     x5 = x.reshape(n, c, 1, 1, w)
     w5 = weight.reshape(k, wc, 1, 1, s)
-    y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, 1, sw), padding=(0, 0, pw), **kwargs)
+    y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, 1, sw), padding=(0, 0, pw), weight_src=weight, **kwargs)
     return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[4])
 
 

tests/kernels/test_conv3d_implicit.py

@@ -254,3 +254,95 @@ 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, :, :].
+# 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__dec_d3_res"),
+    pytest.param(96, 192, 512, 512, id="enc_e1_res1"),
+    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(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)
+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)
+
+
+@_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)