[Kernel] fp8 conv3d: 8-wave GEMM pipeline + BIG_IN fix (#860)

* conv3d fp8: port kernel onto the fp8_gemm 8-wave pipeline

Rewrites conv3d_implicit_fp8 on top of the fp8_gemm_8wave pipeline
(G2SLoader / S2RLoader / Mfma16x16x128), taking FP8 E4M3FN inputs
directly instead of quantizing bf16 in the kernel. Adds WGM L2-swizzle
(configurable + autotuned), generalizes to N/K-partial and tiny-K
shapes via OOB masking, dispatches 1D/2D/3D by filter rank, and fixes
NaN / memory faults on >2 GB tensors by rebasing the buffer descriptor.

Renames conv3d_implicit_autotune.py -> conv3d_autotune.py: the module
now serves both the bf16 (conv3d_implicit.py) and fp8 kernels, so
"implicit" in the name is no longer accurate. Pure rename, no content
change; the two importers are updated.

* conv3d: accept any channel count and unaligned spatial extents

C values that are not a multiple of the gather's vector width (8 bf16 / 16 fp8)
were rejected outright, and the NCHW->NHWC transpose fell back to torch for any
spatial extent that was not a multiple of the same width.

Channels: zero-pad C instead of asserting. The padded channels multiply against
zero weights, so the result is unchanged; rel_l2 against an fp32 reference is
1.66e-3, matching torch's own bf16 error. The weight packers pad too, keyed on
the caller's unpadded tensor so their identity caches still hit.

Spatial extent: the transpose guards were stricter than the kernels need. For
bf16, s has no alignment requirement at all -- the read's tail lands in LDS
columns the write-back guard never visits. For fp8 the read indexes the input
in dwords, so the requirement is s%4 rather than s%16. Only the write side
needs channel alignment, since it stores a full vector along C.

Both transposes leave the input descriptor at max_size on purpose: an exact
num_records makes the hardware zero a 16-byte access that straddles the end of
the tensor, including the valid elements inside it.

bf16 end-to-end vs torch/MIOpen on MI355X (autotuned, 5 trials averaged): the
previously-failing C=3 shape runs at 1.49x, and the stride-2 shapes go 1.04x ->
1.91x, 0.92x -> 1.51x, 0.81x -> 1.36x. Geomean over the real conv shapes
improves 1.295x -> 1.431x with no regression on aligned shapes.

Tests: 57 bf16 (was 30) and 36 fp8 (was 29), covering C in {1,3,4,5,6,12,24,48}
and unaligned spatial extents.

* fix buffer_ops import

* conv3d fp8: migrate raw pointer construction to the fx.* API

Replace the buffer_ops pointer helpers with the current DSL surface:

- extract_base_index(t) -> fx.ptrtoint(fx.get_iter(t))
- create_llvm_ptr(addr, address_space=N) -> <fx ptr>.llvm_ptr, which
  resolves the LLVM address space via the compile backend instead of
  hardcoding <1>/<3>.

The two global raw-address sites share a new _global_ptr_from_addr()
helper, which _make_fp8_buffer_tensor_from_addr() now also uses for its
inttoptr step. The LDS store reuses the recast_iter(u8, ...) idiom the
neighbouring load already used.

Behaviour-preserving. tests/kernels/test_conv3d_implicit_fp8.py 35/35
pass on gfx950; the untested BIG_IN and BIG_OUT branches were both
force-compiled to cover the const_expr paths the shape-gated tests skip.

Co-Authored-By: Claude <noreply@anthropic.com>

* conv3d fp8: migrate buffer_ops resources/loads/stores to fx.copy

Replace the last buffer_ops users with the layout + copy-atom surface, so
the module no longer imports kernels.common.buffer_ops:

- create_buffer_resource -> fx.rocdl.make_buffer_tensor + logical_divide
- buffer_load/buffer_store -> fx.copy through an rmem staging tensor,
  following the StoreC pattern in kernels/gemm/fp8_gemm_utils.py.

Two details worth flagging:

- The transpose input load moves off the i32-element offset (the old
  '// 4') to a u8 buffer tensor, where the copy-atom element offset is
  the byte offset directly. Same address, one less scaling step to get
  wrong.
- The masked epilogue store now routes masked-off lanes to element
  npq*k (one past the end) instead of relying on buffer_store's
  mask=, matching StoreC's oob idiom. The hardware num_records bound
  drops them. This branch only compiles when npq*k*2 <= 2^31, so the
  offset always fits in i32.

_vec_store moves up beside BIG_IN/BIG_OUT because the copy atoms are now
built at kernel entry.

Verified on gfx950 (MI355X):
- test_conv3d_implicit_fp8.py 35/35 pass, re-run with
  FLYDSL_RUNTIME_ENABLE_CACHE=0 to defeat the JIT cache
- BIG_IN / BIG_OUT / TR_BIG / has_bias force-compiled, since all four are
  const_expr branches the default test tier never reaches
- bias has no test coverage at all, so it was checked numerically
  against torch conv3d on 3 shapes (rel_err 1.7e-03, matching no-bias)
- perf unchanged: 2158 vs 2171 TF (1x3x3) and 2489 vs 2496 TF (3x3x3),
  median of 10, within run-to-run noise

Co-Authored-By: Claude <noreply@anthropic.com>

* conv3d bf16: migrate buffer_ops to the fx.* API

Mirrors the fp8 migration in the two preceding commits, so both conv
kernels now use the same surface and neither imports
kernels.common.buffer_ops:

- create_buffer_resource            -> fx.rocdl.make_buffer_tensor
- create_buffer_resource_from_addr  -> make_buffer_ptr(fx.inttoptr(...)),
  reusing the _x_div_from_addr idiom already in this file
- buffer_load / buffer_store        -> fx.copy via an rmem staging tensor
- extract_base_index                -> fx.ptrtoint(fx.get_iter(t))
- create_llvm_ptr                   -> .llvm_ptr

Descriptor semantics are unchanged: every migrated site kept its
num_records (all were the max_size 0xFFFFFFFF default, including the
rebased transpose resources, whose comment explains why an exact bound
would zero the straddling tail read).

The split-K epilogue still needs a raw !llvm.ptr<8> because
buffer_atomic_add takes a resource rather than a tensor, so it derives
one with fx.rocdl.get_buffer_rsrc -- the same approach
kernels/moe/moe_gemm_2stage/gemm2.py uses for its scatter atomics.

_row_chk / _need_chk / _vec_store move up to the enclosing scope since
the copy atoms they select are now built at kernel entry.

Verified on gfx950 (MI355X):
- test_conv3d_implicit.py 57/57 pass with FLYDSL_RUNTIME_ENABLE_CACHE=0
- split-K has no test coverage and is the one path that changed shape
  here, so it was checked numerically against torch conv3d at splitk=4
  and 8, with and without bias (rel_err 6.9e-07..3.1e-03)
- BIG_IN_N1 / BIG_IN_NM / BIG_OUT / splitk / transpose-BIG force-compiled
- transpose verified bit-exact vs torch.permute, including the BIG
  branch on a real 4.1 GiB tensor (2.2e9 elements, >2^31)
- perf unchanged: 723 vs 717 TF (3x3x3) and 657 vs 660 TF (1x3x3)

Co-Authored-By: Claude <noreply@anthropic.com>

* conv3d: tighten migration comments

Drop the comment that restated the buffer-tensor construction and move
the byte-offset note down to the offset arithmetic it explains.

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: jiacao-amd <jiacao@amd.com>
Co-authored-by: Nikolai Protasov <nprotaso@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 4 people

kernels/conv/conv3d_implicit_autotune.py kernels/conv/conv3d_autotune.pykernels/conv/conv3d_implicit_autotune.py renamed to kernels/conv/conv3d_autotune.py

@@ -19,7 +19,6 @@
 from flydsl.expr import arith, const_expr, range_constexpr, rocdl
 from flydsl.expr.rocdl.universal import make_buffer_ptr
 from flydsl.expr.typing import T
-from kernels.common import buffer_ops
 from kernels.common.mem_ops import buffer_atomic_add
 
 TILE_K = 32
@@ -46,12 +45,18 @@ def _autotune_enabled():
 _WEIGHT_CACHE = {}
 
 
+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)
     ent = _WEIGHT_CACHE.get(key)
     if ent is not None and ent[0]() is w:
         return ent[1]
-    wk = w.permute(0, 2, 3, 4, 1).contiguous().reshape(k, kt * kh * kw * c)
+    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)
     return wk
 
@@ -67,16 +72,26 @@ def _prep_weight(w, k, kt, kh, kw, c):
 
 @functools.lru_cache(maxsize=64)
 def compile_transpose_ncdhw_ndhwc(n, c, s):
-    """Transpose flat (N, C, S) -> (N, S, C) (S == T*H*W). Requires c%8==0, s%8==0."""
+    """Transpose flat (N, C, S) -> (N, S, C) (S == T*H*W). Requires c%8==0."""
     grid_s = (s + TR_TILE - 1) // TR_TILE
     grid_c = (c + TR_TILE - 1) // TR_TILE
     elem_ty = fx.BFloat16
     BIG = (n * c * s) > 0x7FFFFFFF
 
+    # 1-D element view so the flat gather/scatter offsets index elements. Both
+    # descriptors stay max_size: an exact num_records would zero the whole
+    # straddling tail read.
+    _TR_FLAT = n * c * s
+    _TR_REBASED_FLAT = 0xFFFFFFFF // BF16_BYTES
+
+    def _flat_div(buf_ptr, elems):
+        return fx.logical_divide(
+            fx.Tensor(fx.make_view(buf_ptr, fx.make_layout(elems, 1))),
+            fx.make_layout(1, 1),
+        )
+
     @flyc.kernel(known_block_size=[TR_THREADS, 1, 1])
     def transpose_kernel(out: fx.Tensor, inp: fx.Tensor):
-        in_rsrc = buffer_ops.create_buffer_resource(inp)
-        out_rsrc = buffer_ops.create_buffer_resource(out)
         lds_alloc = fx.SharedAllocator(static=False)
         lds = lds_alloc.allocate(fx.Array[elem_ty, TR_TILE * _TR_LDS_S, 16]).peek()
 
@@ -88,19 +103,27 @@ class BF16Ty:
         c0 = fx.block_idx.y * TR_TILE
         nb = fx.block_idx.z
         if const_expr(BIG):
+            # Rebase onto this block's tile origin so the per-tile offsets stay in i32.
+            GPtrTy = fx.PointerType.get(elem_ty.ir_type, 1, BF16_BYTES)
+
+            def _rebased(tensor, base_elem):
+                addr = fx.Int64(fx.ptrtoint(fx.get_iter(tensor))) + fx.Int64(base_elem) * fx.Int64(BF16_BYTES)
+                return _flat_div(make_buffer_ptr(fx.inttoptr(GPtrTy, addr)), _TR_REBASED_FLAT)
+
             in_base_elem = fx.Index(nb) * fx.Index(c) * fx.Index(s) + fx.Index(c0) * fx.Index(s) + fx.Index(s0)
-            in_addr = fx.Int64(buffer_ops.extract_base_index(inp)) + fx.Int64(in_base_elem) * fx.Int64(2)
-            in_rsrc = buffer_ops.create_buffer_resource_from_addr(in_addr)
+            in_div = _rebased(inp, in_base_elem)
             out_base_elem = fx.Index(nb) * fx.Index(s) * fx.Index(c) + fx.Index(s0) * fx.Index(c) + fx.Index(c0)
-            out_addr = fx.Int64(buffer_ops.extract_base_index(out)) + fx.Int64(out_base_elem) * fx.Int64(2)
-            out_rsrc = buffer_ops.create_buffer_resource_from_addr(out_addr)
+            out_div = _rebased(out, out_base_elem)
         else:
             in_base = nb * c * s
             out_base = nb * s * c
+            in_div = _flat_div(fx.get_iter(fx.rocdl.make_buffer_tensor(inp)), _TR_FLAT)
+            out_div = _flat_div(fx.get_iter(fx.rocdl.make_buffer_tensor(out)), _TR_FLAT)
+        tr_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_ty)
+        tr_reg = fx.make_rmem_tensor(TR_VEC, elem_ty)
 
         def lds_store_vec8(elem_offset, value):
-            base = fx.Int64(fx.ptrtoint(lds.ptr)) + fx.Int64(elem_offset * 2)
-            ptr = buffer_ops.create_llvm_ptr(base, address_space=3)
+            ptr = (fx.recast_iter(fx.Uint8, lds.ptr) + fx.Int32(elem_offset * 2)).llvm_ptr
             llvm.StoreOp(value, ptr, alignment=16)
 
         def lds_load_scalar(elem_offset):
@@ -120,8 +143,9 @@ def lds_load_scalar(elem_offset):
             else:
                 g = fx.Int32(in_base + cc * s + ss)
             safe = arith.select(valid, g, fx.Int32(0))
-            v = buffer_ops.buffer_load(in_rsrc, safe, vec_width=TR_VEC, dtype=elem_ty)
-            lds_store_vec8(rc * _TR_LDS_S + sv, v)
+            fx.copy(tr_atom, fx.slice(in_div, (None, safe)), tr_reg)
+            v = fx.memref_load_vec(tr_reg)
+            lds_store_vec8(rc * _TR_LDS_S + sv, v.ir_value() if hasattr(v, "ir_value") else v)
 
         llvm.InlineAsmOp(None, [], "s_waitcnt lgkmcnt(0)\n\ts_barrier", "", has_side_effects=True)
 
@@ -139,7 +163,8 @@ def lds_load_scalar(elem_offset):
                     go = fx.Int32(rs * c + cv)
                 else:
                     go = fx.Int32(out_base + ss * c + cc)
-                buffer_ops.buffer_store(vv, out_rsrc, go)
+                fx.memref_store_vec(vv, tr_reg)
+                fx.copy(tr_atom, tr_reg, fx.slice(out_div, (None, go)))
 
     @flyc.jit
     def launch_transpose(out: fx.Tensor, inp: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
@@ -156,7 +181,7 @@ def _ncdhw_to_ndhwc(x, stream):
     """Fast NCDHW->NDHWC via the tiled transpose kernel; falls back to torch."""
     n, c, t, h, w = x.shape
     s = t * h * w
-    if not (x.is_contiguous() and x.dtype == torch.bfloat16 and c % 8 == 0 and s % 8 == 0):
+    if not (x.is_contiguous() and x.dtype == torch.bfloat16 and c % TR_VEC == 0):
         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)
@@ -186,7 +211,7 @@ def compile_conv3d_implicit(
     assert (TILE_M * TILE_K) % BLOCK_VECS == 0, f"A tile {TILE_M}x{TILE_K} not a multiple of {BLOCK_VECS} vecs"
     assert (TILE_N * TILE_K) % BLOCK_VECS == 0, f"B tile {TILE_N}x{TILE_K} not a multiple of {BLOCK_VECS} vecs"
     assert LDG_A_COUNT >= 1 and LDG_B_COUNT >= 1
-    assert c % LDG_VEC == 0
+    assert c % LDG_VEC == 0, f"c={c} must be a multiple of LDG_VEC={LDG_VEC}; use _conv3d_impl to pad"
     assert BLOCK_THREADS <= 1024, f"BLOCK_THREADS={BLOCK_THREADS} exceeds 1024"
 
     do = (d + 2 * pt - kt) // st + 1
@@ -214,10 +239,13 @@ def compile_conv3d_implicit(
 
     n_tail = k % TILE_N != 0
     grid_n = (k + TILE_N - 1) // TILE_N
+    _row_chk = npq % TILE_M != 0
+    _need_chk = _row_chk or n_tail
 
     splitk = max(1, min(splitk, k_tiles))
     tiles_per_split = k_tiles // splitk
     use_splitk = splitk > 1
+    _vec_store = (n == 1) and (not use_splitk) and (dhw % MFMA_C_VALUES == 0) and (not BIG_OUT)
 
     # Software-pipeline depth. 4 stages is optimal across all shapes on gfx950 --
     # even short-K, memory-bound 3x1x1 depends more (not less) on deep prefetch to
@@ -246,7 +274,20 @@ def compile_conv3d_implicit(
 
     @flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1])
     def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor):
-        y_rsrc = buffer_ops.create_buffer_resource(y)
+        y_buf = fx.rocdl.make_buffer_tensor(y)
+        if const_expr(use_splitk):
+            # buffer_atomic_add needs the raw !llvm.ptr<8> descriptor, not a tensor.
+            y_rsrc = fx.rocdl.get_buffer_rsrc(fx.get_iter(y_buf))
+        else:
+            y_div = fx.logical_divide(
+                fx.Tensor(fx.make_view(fx.get_iter(y_buf), fx.make_layout(npq * k, 1))),
+                fx.make_layout(1, 1),
+            )
+            y_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), elem_ty)
+            y_reg_1 = fx.make_rmem_tensor(1, elem_ty)
+            if const_expr(_vec_store):
+                y_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), elem_ty)
+                y_reg_4 = fx.make_rmem_tensor(MFMA_C_VALUES, elem_ty)
         # Buffer tensors for the im2col gather, flattened to a 1-D element view so the
         # per-thread flat gather offset indexes elements (multi-dim views would not).
         w_buf0 = fx.rocdl.make_buffer_tensor(weight, max_size=False)
@@ -257,7 +298,9 @@ def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias:
             x_buf = fx.Tensor(fx.make_view(fx.get_iter(x_buf0), fx.make_layout(n * c * d * h * w, 1)))
             x_div = fx.logical_divide(x_buf, fx.make_layout(1, 1))
         if const_expr(has_bias):
-            bias_rsrc = buffer_ops.create_buffer_resource(bias)
+            bias_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(bias), fx.make_layout(1, 1))
+            bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32)
+            bias_reg = fx.make_rmem_tensor(1, fx.Float32)
 
         lds_alloc = fx.SharedAllocator(static=False)
         a_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_A_SIZE, 16]).peek()
@@ -285,6 +328,7 @@ def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias:
         # BIG_IN (>2GB): flat buffer tensor from a rebased address with explicit ~2GB
         # num_records (same mechanism as x_div, replacing create_buffer_resource_from_addr).
         GXPtrTy = fx.PointerType.get(elem_ty.ir_type, 1, BF16_BYTES) if const_expr(BIG_IN) else None
+        GYPtrTy = fx.PointerType.get(elem_ty.ir_type, 1, BF16_BYTES) if const_expr(BIG_OUT) else None
 
         def _x_div_from_addr(addr_i64):
             gptr = fx.inttoptr(GXPtrTy, addr_i64)
@@ -299,10 +343,10 @@ def _x_div_from_addr(addr_i64):
             base_t = ot_base0 - fx.Index(pt)
             base_t = arith.select(base_t < fx.Index(0), fx.Index(0), base_t)
             x_base_elem = ((nbase * fx.Index(d) + base_t) * fx.Index(h) + fx.Index(0)) * fx.Index(w) * fx.Index(c)
-            x_addr = fx.Int64(buffer_ops.extract_base_index(x)) + fx.Int64(x_base_elem) * fx.Int64(2)
+            x_addr = fx.Int64(fx.ptrtoint(fx.get_iter(x))) + fx.Int64(x_base_elem) * fx.Int64(2)
             x_div_big = _x_div_from_addr(x_addr)
         if const_expr(BIG_IN_NM):
-            x_base_addr = fx.Int64(buffer_ops.extract_base_index(x))
+            x_base_addr = fx.Int64(fx.ptrtoint(fx.get_iter(x)))
 
         wid = tid // WARP_SIZE
         lane = tid % WARP_SIZE
@@ -551,16 +595,12 @@ def do_compute(acc_values, a_frag_values, b_frag_values):
                 rocdl.sched_vmem(LDG_A_COUNT + LDG_B_COUNT)
             acc = do_compute(acc, a_frags, b_frags)
 
-        _row_chk = npq % TILE_M != 0
-        _need_chk = _row_chk or n_tail
-        _vec_store = (n == 1) and (not use_splitk) and (dhw % MFMA_C_VALUES == 0) and (not BIG_OUT)
-
         if const_expr(BIG_OUT):
-            y_elem_base = fx.Int64(buffer_ops.extract_base_index(y))
+            y_elem_base = fx.Int64(fx.ptrtoint(fx.get_iter(y)))
 
         def _big_store(off_nk_i64, value):
             addr = y_elem_base + off_nk_i64 * fx.Int64(BF16_BYTES)
-            ptr = buffer_ops.create_llvm_ptr(addr, address_space=1)
+            ptr = fx.inttoptr(GYPtrTy, addr).llvm_ptr
             llvm.StoreOp(value.ir_value() if hasattr(value, "ir_value") else value, ptr, alignment=2)
 
         def _valid_raw(row, col):
@@ -582,7 +622,8 @@ def store_acc():
                         col_i = fx.Int32(col)
                         if const_expr(n_tail):
                             col_i = arith.select(col < fx.Index(k), col_i, fx.Int32(0))
-                        bias_val = fx.Float32(buffer_ops.buffer_load(bias_rsrc, col_i, vec_width=1, dtype=fx.Float32))
+                        fx.copy(bias_atom, fx.slice(bias_div, (None, col_i)), bias_reg)
+                        bias_val = fx.Float32(fx.memref_load_vec(bias_reg)[0])
 
                     if const_expr(_vec_store):
                         row0 = fx.Index(row_base)
@@ -594,7 +635,8 @@ def _emit_vec():
                                 cval = (a[i] + bias_val) if const_expr(has_bias) else a[i]
                                 vals.append(cval.to(elem_ty))
                             v4 = fx.Vector.from_elements(vals, dtype=elem_ty)
-                            buffer_ops.buffer_store(v4, y_rsrc, off_nk0)
+                            fx.memref_store_vec(v4, y_reg_4)
+                            fx.copy(y_atom_4, y_reg_4, fx.slice(y_div, (None, fx.Int32(off_nk0))))
 
                         if const_expr(_need_chk):
                             if _valid_raw(row0, col):
@@ -624,7 +666,8 @@ def _emit():
                                 if const_expr(BIG_OUT):
                                     _big_store(fx.Int64(off_nk), cval)
                                 else:
-                                    buffer_ops.buffer_store(cval, y_rsrc, off_nk)
+                                    fx.memref_store_vec(fx.Vector.filled(1, cval, elem_ty), y_reg_1)
+                                    fx.copy(y_atom_1, y_reg_1, fx.slice(y_div, (None, fx.Int32(off_nk))))
 
                         if const_expr(_need_chk):
                             if _valid_raw(row, col):
@@ -696,14 +739,20 @@ def _conv3d_impl(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=
     ho = (h + 2 * ph - kh) // sh + 1
     wo = (w + 2 * pw - kw) // sw + 1
     npq = n * do * ho * wo
+
+    # Zero-pad C to the gather's vector width; padded channels see zero weights.
+    cp = _pad_channels(c)
+    if cp != c:
+        x = torch.nn.functional.pad(x, (0, 0, 0, 0, 0, 0, 0, cp - c))
+    c = cp
     crs = c * kt * kh * kw
 
     launch_stream = torch.cuda.current_stream() if stream is None else stream
     has_bias = bias is not None
     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, c)
+    w_packed = _prep_weight(weight, k, kt, kh, kw, wc)
 
     shape = (n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias)
 
@@ -723,7 +772,7 @@ def _run(the_tile, the_wgm=1):
         chosen_tile = tuple(tile)
         chosen_wgm = 1
     elif autotune or (autotune is None and _autotune_enabled()):
-        from kernels.conv.conv3d_implicit_autotune import BF16_CANDIDATES, WGM_VALUES, autotune_conv3d
+        from kernels.conv.conv3d_autotune import BF16_CANDIDATES, WGM_VALUES, autotune_conv3d
 
         candidates = [(t, w) for t in BF16_CANDIDATES for w in WGM_VALUES]
         best = autotune_conv3d("bf16", shape, "bf16", candidates, x.device, lambda tw: _run(tw[0], tw[1])[0])