Keep the wide N tile when K per group only partly fills it.

_pick_tile rejected every tile with TILE_N > K/groups, so the Qwen-Image VAE
layers with cout=96 dropped from (128,128,2,4) to (64,64,2,2) and lost 30-47%
of their GEMM time: the narrower tile halves the waves per block and doubles
the A traffic per output element, which costs far more than the 25% of masked
N columns. Accept a tile whose N span is at least TILE_MIN_N_FILL full; the
wave-count check still demotes it when the problem cannot fill the device.

Measured on gfx950, bf16, 3x3, wall clock (independent process per shape, new
and old selector alternated in-process, min of 5):

  shape             calls   before    after
  96->96  @1024x1024   10   0.527    0.394
  96->96  @1328x1328   10   0.973    0.692
  192->96 @1024x1024    1   0.925    0.694
  96->96  @1025 s2      1   0.269    0.224
  3->96   @1024x1024    1   0.142    0.130

The other 13 VAE shapes keep their tile and are unchanged. Weighted over the
VAE call mix: 1024 path 14.46 -> 12.82 ms (1.29x -> 1.45x vs MIOpen), the two
1328 layers 12.04 -> 9.22 ms (1.37x -> 1.79x).

Author: amd-nprotaso

kernels/conv/conv3d_implicit.py

@@ -53,6 +53,8 @@
 
 TILE_MIN_WAVES_PER_CU = 6
 
+TILE_MIN_N_FILL = 0.75
+
 PADDING_MODES = ("zeros", "reflect", "replicate", "circular")
 
 CONV_COMPILE_HINTS = {}
@@ -954,7 +956,11 @@ def _num_cu(device):
 
 def _pick_tile(npq, k, groups, device):
     kg = k // groups
-    legal = [t for t in TILE_LADDER if t[1] <= kg] or [TILE_LADDER[-1]]
+    # A tile wider than kg is still worth its masked columns: it keeps more waves per
+    # block and halves the A traffic per output element. Below TILE_MIN_N_FILL the
+    # wasted columns take over; the wave-count check below demotes it again when the
+    # problem is too small to fill the device.
+    legal = [t for t in TILE_LADDER if kg >= t[1] * TILE_MIN_N_FILL] or [TILE_LADDER[-1]]
     target = TILE_MIN_WAVES_PER_CU * _num_cu(device)
     for tile_m, tile_n, wave_m, wave_n in legal:
         blocks = ((npq + tile_m - 1) // tile_m) * groups * ((kg + tile_n - 1) // tile_n)

tests/kernels/test_conv3d_implicit.py

@@ -15,7 +15,7 @@
 import torch.nn.functional as F
 
 from flydsl.runtime.device import get_rocm_arch
-from kernels.conv.conv3d_implicit import conv3d_implicit
+from kernels.conv.conv3d_implicit import _pick_tile, conv3d_implicit
 
 pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower]
 
@@ -129,6 +129,26 @@ def test_conv3d_tile_configs(tile):
     assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2)
 
 
+@_skip_non_cdna4
+@pytest.mark.parametrize("k,groups", [(96, 1), (192, 2)])
+def test_conv2d_auto_tile_n_tail(k, groups):
+    """K per group between TILE_MIN_N_FILL*TILE_N and TILE_N must keep the wide tile."""
+    torch.manual_seed(4400 + k)
+    c, h, w = 64 * groups, 256, 256
+    x = torch.randn((1, c, h, w), device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn((k, c // groups, 3, 3), device="cuda", dtype=torch.bfloat16)
+    bias = torch.randn((k,), device="cuda", dtype=torch.float32)
+
+    assert _pick_tile(h * w, k, groups, x.device)[1] == 128
+
+    y = conv3d_implicit(x, weight, bias=bias, padding=1, groups=groups)
+    y_ref = F.conv2d(x, weight, bias=bias.to(torch.bfloat16), padding=1, groups=groups)
+    torch.cuda.synchronize()
+
+    assert y.shape == y_ref.shape
+    assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2)
+
+
 @_skip_non_cdna4
 def test_conv3d_autotune(tmp_path, monkeypatch):
     monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "at"))