diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py
index 96dd73a6ed..52b6779ebe 100644
--- a/kernels/attention/flash_attn_utils.py
+++ b/kernels/attention/flash_attn_utils.py
@@ -2596,8 +2596,8 @@ def coop_dma_v_nomajor(self, tile_start, buf_id=0):
         traits = ctx.traits
         pid = self._tile_page_id(tile_start)
         paddr = as_mlir_value(self._v_dma_base_i64 + fx.Int64(pid) * self._v_dma_page_bytes)
-        rsrc = buffer_ops.create_buffer_resource_from_addr(
-            paddr, num_records_bytes=as_mlir_value(self._v_dma_page_bytes)
+        rsrc = fx.rocdl.get_buffer_rsrc(
+            buffer_ops.create_buffer_resource_from_addr(paddr, num_records_bytes=as_mlir_value(self._v_dma_page_bytes))
         )
         if const_expr(isinstance(buf_id, int)):
             vb = self._v_dma_lds_base + fx.Index((traits.LDS_V_BASE + buf_id * traits.LDS_V_TILE_SIZE) * 2)
@@ -2630,14 +2630,18 @@ def init_dma(self):
         self._dma_soff = fx.Int32(0)
         self._dma_off = fx.Int32(0)
         self._dma_aux = fx.Int32(1)
-        self.k_rsrc = buffer_ops.create_buffer_resource(
-            ctx.K, max_size=False, num_records_bytes=ctx.kv_nrec_bytes, base_byte_offset=ctx.kv_batch_byte_off
+        self.k_rsrc = fx.rocdl.get_buffer_rsrc(
+            buffer_ops.create_buffer_resource(
+                ctx.K, max_size=False, num_records_bytes=ctx.kv_nrec_bytes, base_byte_offset=ctx.kv_batch_byte_off
+            )
         )
         self.NUM_DMA_K = (traits.BLOCK_N * traits.K_STRIDE * 2) // self.DMA_BATCH_BYTES
         self.LANES_PER_K_ROW = traits.HEAD_DIM * 2 // self.DMA_BYTES
         self.ROWS_PER_DMA_BATCH = self.DMA_BATCH_BYTES // (traits.HEAD_DIM * 2)
-        self.v_rsrc = buffer_ops.create_buffer_resource(
-            ctx.V, max_size=False, num_records_bytes=ctx.kv_nrec_bytes, base_byte_offset=ctx.kv_batch_byte_off
+        self.v_rsrc = fx.rocdl.get_buffer_rsrc(
+            buffer_ops.create_buffer_resource(
+                ctx.V, max_size=False, num_records_bytes=ctx.kv_nrec_bytes, base_byte_offset=ctx.kv_batch_byte_off
+            )
         )
         self.NUM_DMA_V = (traits.BLOCK_N * traits.V_STRIDE * 2) // self.DMA_BATCH_BYTES
         self.LANES_PER_V_ROW = traits.HEAD_DIM * 2 // self.DMA_BYTES
diff --git a/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py b/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py
index f382f6b657..13129cc4cd 100644
--- a/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py
+++ b/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py
@@ -365,7 +365,7 @@ def _vt_perm(src_hi, src_lo, sel):
 
     # ---- Buffer resources ----
     query_rsrc = buffer_ops.create_buffer_resource(query)
-    kv_rsrc = buffer_ops.create_buffer_resource(kv_buffer)
+    kv_rsrc = fx.rocdl.get_buffer_rsrc(buffer_ops.create_buffer_resource(kv_buffer))
     kv_page_indices_rsrc = buffer_ops.create_buffer_resource(kv_page_indices)
     work_indptr_rsrc = buffer_ops.create_buffer_resource(work_indptr)
     work_info_set_rsrc = buffer_ops.create_buffer_resource(work_info_set)
diff --git a/kernels/common/buffer_ops.py b/kernels/common/buffer_ops.py
index 6f725c03c4..f843c77fa3 100644
--- a/kernels/common/buffer_ops.py
+++ b/kernels/common/buffer_ops.py
@@ -6,6 +6,8 @@
 This module provides high-level Python wrappers for AMD CDNA3/CDNA4 buffer operations.
 Buffer operations use a scalar base pointer and per-thread offsets for efficient memory access.
 
+``create_buffer_resource`` returns a byte-addressed buffer-descriptor pointer
+
 Example:
     >>> from kernels.common import buffer_ops
     >>> from flydsl.expr import arith
@@ -24,11 +26,12 @@
     >>> buffer_ops.buffer_store(data, rsrc, offset)
 """
 
-from typing import Optional, Union
+from __future__ import annotations
 
+import flydsl.expr as fx
 from flydsl._mlir import ir
 from flydsl._mlir.dialects import arith as std_arith
-from flydsl._mlir.dialects import llvm, rocdl
+from flydsl._mlir.dialects import llvm
 from flydsl._mlir.extras import types as T
 from flydsl.expr.meta import dsl_loc_tracing
 from flydsl.runtime.device import is_rdna_arch
@@ -37,6 +40,8 @@
 def _get_buffer_flags(arch=None):
     """Get AMD buffer resource descriptor (V#) flags word (bits 127:96).
 
+     ``flash_attn_utils`` builds a couple of descriptors by hand and still calls it.
+
     Constructs the 32-bit flags field for rocdl.make.buffer.rsrc, following the
     same logic as LLVM's AMDGPUToROCDL makeBufferRsrc():
       https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -72,16 +77,18 @@ def _get_buffer_flags(arch=None):
 
 
 __all__ = [
-    "create_llvm_ptr",
-    "get_element_ptr",
-    "create_buffer_resource",
-    "create_buffer_resource_from_addr",
+    "BufferResourceDescriptor",
     "buffer_load",
     "buffer_store",
-    "BufferResourceDescriptor",
+    "create_buffer_resource",
+    "create_buffer_resource_from_addr",
+    "create_llvm_ptr",
     "extract_base_index",
+    "get_element_ptr",
 ]
 
+_MAX_NUM_RECORDS = 0xFFFFFFFF
+
 
 def _unwrap_value(value):
     """Recursively unwrap ArithValue or similar wrappers to get the actual MLIR value.
@@ -90,6 +97,7 @@ def _unwrap_value(value):
     - FlyDSL ArithValue (has ._value)
     - flyc DSL Numeric like fx.Int32 (has .ir_value() method)
     - flyc ArithValue (is already ir.Value subclass)
+    - a bare single-result OpView (has .result), e.g. ``arith.TruncIOp(...)``
     """
     # DSL Numeric (Int32, Float32, etc.) — use ir_value() to materialize
     if hasattr(value, "ir_value") and not isinstance(value, ir.Value):
@@ -99,6 +107,8 @@ def _unwrap_value(value):
     while depth < max_depth and not isinstance(value, ir.Value):
         if hasattr(value, "_value"):
             value = value._value
+        elif hasattr(value, "result"):
+            value = value.result
         elif hasattr(value, "value"):
             value = value.value
         else:
@@ -118,24 +128,6 @@ def _create_i32_constant(value: int) -> ir.Value:
     return _unwrap_value(op.result)
 
 
-@dsl_loc_tracing
-def _create_i16_constant(value: int) -> ir.Value:
-    """Create i16 constant using standard MLIR arith dialect."""
-    i16_type = T.i16()
-    attr = ir.IntegerAttr.get(i16_type, value)
-    op = std_arith.ConstantOp(i16_type, attr)
-    return _unwrap_value(op.result)
-
-
-@dsl_loc_tracing
-def _create_i64_constant(value: int) -> ir.Value:
-    """Create i64 constant using standard MLIR arith dialect."""
-    i64_type = T.i64()
-    attr = ir.IntegerAttr.get(i64_type, value)
-    op = std_arith.ConstantOp(i64_type, attr)
-    return _unwrap_value(op.result)
-
-
 @dsl_loc_tracing
 def _ptr8_to_v4i32(ptr8_val) -> ir.Value:
     """Reinterpret a buffer resource (!llvm.ptr<8>) as a <4 x i32> vector.
@@ -149,13 +141,41 @@ def _ptr8_to_v4i32(ptr8_val) -> ir.Value:
     return llvm.bitcast(v4i32_ty, i128_val)
 
 
+def _as_num_records(num_records_bytes) -> fx.Int64 | None:
+    """Normalize a descriptor byte count to ``fx.Int64`` for ``make_buffer_ptr``."""
+    if num_records_bytes is None:
+        return None
+    if isinstance(num_records_bytes, int):
+        num_records_bytes = max(0, min(int(num_records_bytes), _MAX_NUM_RECORDS))
+    return fx.Int64(num_records_bytes)
+
+
+@dsl_loc_tracing
+def _byte_buffer_ptr(global_ptr, num_records_bytes, base_byte_offset=None):
+    """Wrap a global pointer in a byte-addressed buffer-resource pointer."""
+    ptr = fx.recast_iter(fx.Int8, global_ptr)
+    if base_byte_offset is not None:
+        if isinstance(base_byte_offset, ir.Value) and not isinstance(base_byte_offset.type, ir.IndexType):
+            base_byte_offset = fx.Int32(base_byte_offset)
+        ptr = ptr + base_byte_offset
+    return fx.rocdl.make_buffer_ptr(ptr, num_records_bytes=_as_num_records(num_records_bytes))
+
+
+@dsl_loc_tracing
+def _add_soffset_bytes(byte_offset: ir.Value, soffset_bytes) -> ir.Value:
+    """Fold ``soffset_bytes`` into the byte offset."""
+    if soffset_bytes is None:
+        return byte_offset
+    soffset = fx.Int32(soffset_bytes).ir_value()
+    return _unwrap_value(std_arith.AddIOp(byte_offset, soffset).result)
+
+
 @dsl_loc_tracing
 def create_llvm_ptr(value, address_space: int = 0) -> ir.Value:
     """Create an LLVM pointer from an integer or index value."""
     value = _unwrap_value(value)
     if isinstance(value.type, ir.IndexType):
-        i64_type = T.i64()
-        value = _unwrap_value(std_arith.IndexCastOp(i64_type, value).result)
+        value = fx.Int64(value).ir_value()
     ptr_type = ir.Type.parse(f"!llvm.ptr<{address_space}>")
     return llvm.IntToPtrOp(ptr_type, value).result
 
@@ -181,15 +201,15 @@ def extract_base_index(tensor, address_space: int = 1) -> ir.Value:
     ptr_type = ir.Type.parse(f"!llvm.ptr<{address_space}>")
     ptr = _fly.extract_aligned_pointer_as_index(ptr_type, raw)
     i64_val = llvm.PtrToIntOp(ir.IntegerType.get_signless(64), ptr).result
-    return _unwrap_value(std_arith.IndexCastOp(ir.IndexType.get(), i64_val).result)
+    return fx.Index(i64_val).ir_value()
 
 
 @dsl_loc_tracing
 def get_element_ptr(
     base_ptr,
-    byte_offset: Union[int, ir.Value, None] = None,
+    byte_offset: int | ir.Value | None = None,
     static_byte_offset: int = 0,
-    elem_type: Optional[ir.Type] = None,
+    elem_type: ir.Type | None = None,
     no_wrap_flags=None,
 ) -> ir.Value:
     """Build an LLVM GEP from a base pointer plus byte offsets."""
@@ -212,8 +232,7 @@ def get_element_ptr(
     else:
         offset_val = _unwrap_value(byte_offset)
         if isinstance(offset_val.type, ir.IndexType):
-            i64_type = T.i64()
-            offset_val = _unwrap_value(std_arith.IndexCastOp(i64_type, offset_val).result)
+            offset_val = fx.Int64(offset_val).ir_value()
         elif not isinstance(offset_val.type, ir.IntegerType):
             raise TypeError("byte_offset must be int, index, or integer-typed MLIR value; " f"got {offset_val.type}")
 
@@ -236,6 +255,28 @@ def get_element_ptr(
     ).result
 
 
+def _num_records_from_memref_type(memref_val) -> int | None:
+    try:
+        mt = ir.MemRefType(_unwrap_value(memref_val).type)
+        shape = list(mt.shape)
+        if any(int(d) < 0 for d in shape):
+            return None
+        # Compute element size in bytes (scalar element type).
+        elem_bits = getattr(mt.element_type, "width", None)
+        if elem_bits is None:
+            return None
+        elem_bytes = int(elem_bits) // 8
+        if elem_bytes <= 0:
+            return None
+        num_elems = 1
+        for d in shape:
+            num_elems *= int(d)
+        return int(num_elems) * int(elem_bytes)
+    # best-effort size probe: any failure just means "unknown"
+    except Exception:  # noqa: BLE001
+        return None
+
+
 class BufferResourceDescriptor:
     """AMD Buffer Resource Descriptor
 
@@ -244,12 +285,10 @@ class BufferResourceDescriptor:
     - stride: Stride for structured buffers (typically 0 for contiguous)
     - num_records: Buffer size in bytes
     - flags: Data format and access flags
-
-    The descriptor is stored in a special LLVM pointer type (!llvm.ptr<8>)
     """
 
-    def __init__(self, rsrc: ir.Value):
-        """Initialize with ROCDL resource descriptor value."""
+    def __init__(self, rsrc):
+        """Initialize with the buffer-descriptor pointer."""
         self.rsrc = rsrc
 
     @staticmethod
@@ -259,9 +298,9 @@ def from_memref(
         stride: int = 0,
         max_size: bool = True,
         data_format: str = "f32",
-        num_records_bytes: Optional[Union[int, ir.Value]] = None,
-        base_byte_offset: Optional[Union[int, ir.Value]] = None,
-    ) -> "BufferResourceDescriptor":
+        num_records_bytes: int | ir.Value | None = None,
+        base_byte_offset: int | ir.Value | None = None,
+    ) -> BufferResourceDescriptor:
         """Create buffer resource descriptor from memref.
 
         Args:
@@ -279,89 +318,22 @@ def from_memref(
         Example:
             >>> rsrc = BufferResourceDescriptor.from_memref(A)
         """
-        # Extract raw pointer from fly.memref.
-        raw_val = _unwrap_value(memref_val)
-        from flydsl._mlir.dialects import fly as _fly
-
-        ptr_type = ir.Type.parse("!llvm.ptr")
-        base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, raw_val)
-        if base_byte_offset is not None:
-            base_ptr = get_element_ptr(base_ptr, byte_offset=base_byte_offset)
-
-        # Create buffer resource descriptor
-        flags_val = _get_buffer_flags()
-        flags = _create_i32_constant(flags_val)
-        stride_val = _create_i16_constant(stride)
-
-        def _num_records_from_memref_type() -> Optional[int]:
-            """Best-effort: derive logical buffer size (in bytes) from static memref type."""
-            try:
-                mt = ir.MemRefType(_unwrap_value(memref_val).type)
-                shape = list(mt.shape)
-                if any(int(d) < 0 for d in shape):
-                    return None
-                # Compute element size in bytes (scalar element type).
-                elem_t = mt.element_type
-                elem_bits = getattr(elem_t, "width", None)
-                if elem_bits is None:
-                    return None
-                elem_bytes = int(elem_bits) // 8
-                if elem_bytes <= 0:
-                    return None
-                num_elems = 1
-                for d in shape:
-                    num_elems *= int(d)
-                return int(num_elems) * int(elem_bytes)
-            except Exception:
-                return None
-
-        if num_records_bytes is not None:
-            # Caller-provided size in BYTES (preferred for exact hardware OOB behavior).
-            if isinstance(num_records_bytes, int):
-                nbytes = int(num_records_bytes)
-                if nbytes <= 0:
-                    nbytes = 0
-                # Descriptor uses i32 bytes; clamp to the max representable.
-                if nbytes > 0xFFFFFFFF:
-                    nbytes = 0xFFFFFFFF
-                num_records = _create_i64_constant(nbytes)
-            else:
-                v = _unwrap_value(num_records_bytes)
-                i64_type = T.i64()
-                if not isinstance(v.type, ir.IntegerType) or v.type.width != 64:
-                    if isinstance(v.type, ir.IndexType):
-                        op = std_arith.IndexCastOp(i64_type, v)
-                    else:
-                        op = std_arith.ExtSIOp(i64_type, v)
-                    v = _unwrap_value(op.result)
-                num_records = v
-        elif max_size:
-            # Use max for flexibility (hardware will check actual bounds)
-            # Note: FlyDSL's rocdl.make.buffer.rsrc requires i32, not i64
-            num_records = _create_i64_constant(0xFFFFFFFF)  # FALLBACK_MAX_SIZE
-        else:
-            # Use the logical memref size (in bytes) for hardware OOB checking.
-            nbytes = _num_records_from_memref_type()
-            if nbytes is None:
-                # Fall back to max-size if we can't infer statically.
-                num_records = _create_i64_constant(0xFFFFFFFF)
-            else:
-                if nbytes > 0xFFFFFFFF:
-                    nbytes = 0xFFFFFFFF
-                num_records = _create_i64_constant(int(nbytes))
-
-        # Create resource descriptor (returns !llvm.ptr<8>)
-        rsrc_type = ir.Type.parse("!llvm.ptr<8>")
-        rsrc = rocdl.MakeBufferRsrcOp(rsrc_type, base_ptr, stride_val, num_records, flags).result
-
-        return BufferResourceDescriptor(rsrc)
+        return BufferResourceDescriptor(
+            create_buffer_resource(
+                memref_val,
+                stride,
+                max_size,
+                num_records_bytes=num_records_bytes,
+                base_byte_offset=base_byte_offset,
+            )
+        )
 
 
 @dsl_loc_tracing
 def create_buffer_resource_from_addr(
     addr_i64: ir.Value,
     *,
-    num_records_bytes: Optional[Union[int, ir.Value]] = None,
+    num_records_bytes: int | ir.Value | None = None,
 ) -> ir.Value:
     """Create AMD buffer resource descriptor from a raw i64 device address.
 
@@ -375,32 +347,15 @@ def create_buffer_resource_from_addr(
         num_records_bytes: Optional buffer size in bytes for hardware OOB checking.
 
     Returns:
-        ROCDL buffer resource descriptor (!llvm.ptr<8>).
+        Byte-addressed buffer-descriptor pointer (``!fly.ptr<i8, BufferDesc>``).
 
     Example:
         >>> rsrc = create_buffer_resource_from_addr(raw_addr_i64)
         >>> data = buffer_load(rsrc, i32_zero, vec_width=4, dtype=T.i32)
     """
-    addr_i64 = _unwrap_value(addr_i64)
-    ptr_type = ir.Type.parse("!llvm.ptr")
-    base_ptr = llvm.IntToPtrOp(ptr_type, addr_i64).result
-    flags = _create_i32_constant(_get_buffer_flags())
-    stride = _create_i16_constant(0)
-    if num_records_bytes is None:
-        num_records = _create_i64_constant(0xFFFFFFFF)
-    elif isinstance(num_records_bytes, int):
-        nbytes = max(0, min(int(num_records_bytes), 0xFFFFFFFF))
-        num_records = _create_i64_constant(nbytes)
-    else:
-        num_records = _unwrap_value(num_records_bytes)
-        i64_type = T.i64()
-        if not isinstance(num_records.type, ir.IntegerType) or num_records.type.width != 64:
-            if isinstance(num_records.type, ir.IndexType):
-                num_records = _unwrap_value(std_arith.IndexCastOp(i64_type, num_records).result)
-            else:
-                num_records = _unwrap_value(std_arith.ExtSIOp(i64_type, num_records).result)
-    rsrc_type = ir.Type.parse("!llvm.ptr<8>")
-    return rocdl.MakeBufferRsrcOp(rsrc_type, base_ptr, stride, num_records, flags).result
+    ptr_ty = fx.PointerType.get(T.i8(), address_space=fx.AddressSpace.Global, alignment=16)
+    base_ptr = fx.inttoptr(ptr_ty, fx.Int64(_unwrap_value(addr_i64)))
+    return _byte_buffer_ptr(base_ptr, num_records_bytes)
 
 
 @dsl_loc_tracing
@@ -409,36 +364,35 @@ def create_buffer_resource(
     stride: int = 0,
     max_size: bool = True,
     *,
-    num_records_bytes: Optional[Union[int, ir.Value]] = None,
-    base_byte_offset: Optional[Union[int, ir.Value]] = None,
+    num_records_bytes: int | ir.Value | None = None,
+    base_byte_offset: int | ir.Value | None = None,
 ) -> ir.Value:
     """Create AMD buffer resource descriptor from memref.
 
-    This is a simplified wrapper around BufferResourceDescriptor.from_memref()
-    that returns the raw ROCDL resource value.
-
     Args:
         memref_val: Memref value
-        stride: Buffer stride (0 for contiguous)
+        stride: Buffer stride (0 for contiguous); only 0 is supported
         max_size: Use maximum buffer size
         num_records_bytes: Override buffer size in bytes.
         base_byte_offset: Optional byte offset added to the descriptor base pointer.
 
     Returns:
-        ROCDL buffer resource descriptor (!llvm.ptr<8>)
+        Byte-addressed buffer-descriptor pointer (``!fly.ptr<i8, BufferDesc>``), for
+        :func:`buffer_load` / :func:`buffer_store`. Pass it through
+        :func:`flydsl.expr.rocdl.get_buffer_rsrc` to recover the raw ``!llvm.ptr<8>``
+        resource that raw ROCDL intrinsics and inline asm expect.
 
     Example:
         >>> rsrc = create_buffer_resource(A)
         >>> data = buffer_load(rsrc, offset)
     """
-    desc = BufferResourceDescriptor.from_memref(
-        memref_val,
-        stride,
-        max_size,
-        num_records_bytes=num_records_bytes,
-        base_byte_offset=base_byte_offset,
-    )
-    return desc.rsrc
+    if stride != 0:
+        raise ValueError(f"create_buffer_resource: only stride=0 (contiguous) is supported, got {stride}")
+
+    if num_records_bytes is None and not max_size:
+        num_records_bytes = _num_records_from_memref_type(memref_val)
+
+    return _byte_buffer_ptr(fx.get_iter(memref_val), num_records_bytes, base_byte_offset)
 
 
 @dsl_loc_tracing
@@ -447,9 +401,9 @@ def buffer_load(
     offset: ir.Value,
     vec_width: int = 4,
     dtype=None,
-    mask: Optional[ir.Value] = None,
+    mask: ir.Value | None = None,
     cache_modifier: int = 0,
-    soffset_bytes: Optional[Union[int, ir.Value]] = None,
+    soffset_bytes: int | ir.Value | None = None,
     is_scalar: bool = False,
 ) -> ir.Value:
     """AMD buffer load operation.
@@ -458,7 +412,7 @@ def buffer_load(
     Uses hardware-level bounds checking and vectorization.
 
     Args:
-        rsrc: Buffer resource descriptor (!llvm.ptr<8>)
+        rsrc: Buffer descriptor pointer from :func:`create_buffer_resource`
         offset: Offset in elements (i32 type)
         vec_width: Vector width (1, 2, or 4)
         dtype: Element data type (None for f32, or ir.F32Type, etc.)
@@ -506,9 +460,7 @@ def buffer_load(
     offset = _unwrap_value(offset)
 
     # Convert offset to i32 if needed
-    if not isinstance(offset.type, ir.IntegerType) or offset.type.width != 32:
-        op = std_arith.IndexCastOp(T.i32(), offset)
-        offset = _unwrap_value(op.result)
+    offset = fx.Int32(offset).ir_value()
 
     # IMPORTANT: Buffer load offset is in BYTES, not elements!
     # For vec4xf32, each element is 4 bytes, so multiply offset by 4
@@ -530,10 +482,11 @@ def buffer_load(
     else:
         result_type = ir.VectorType.get([vec_width], dtype)
 
-    # Scalar/uniform load path: emit s.buffer.load with a v4i32 resource and the
-    # byte offset computed above. Returns i32 (vec_width 1) or v4i32 (vec_width 4).
+    # Scalar/uniform load path: s.buffer.load is an SMEM instruction with no CopyOp type,
+    # so it stays on the raw intrinsic and needs the raw v4i32 resource. Returns i32
+    # (vec_width 1) or v4i32 (vec_width 4).
     if is_scalar:
-        rsrc_v4 = _ptr8_to_v4i32(rsrc)
+        rsrc_v4 = _ptr8_to_v4i32(fx.rocdl.get_buffer_rsrc(rsrc))
         cache_policy = _create_i32_constant(cache_modifier)
         suffix = "i32" if vec_width == 1 else "v4i32"
         return llvm.call_intrinsic(
@@ -544,23 +497,16 @@ def buffer_load(
             [],
         )
 
-    # Create instruction offset and aux flags
-    if soffset_bytes is None:
-        soffset = _create_i32_constant(0)
-    else:
-        if isinstance(soffset_bytes, int):
-            soffset = _create_i32_constant(soffset_bytes)
-        else:
-            soffset = _unwrap_value(soffset_bytes)
-            if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32:
-                op = std_arith.IndexCastOp(T.i32(), soffset)
-                soffset = _unwrap_value(op.result)
-    aux_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) if cache_modifier else None
-
-    # Emit buffer load
-    load_op = rocdl.RawPtrBufferLoadOp(result_type, rsrc, offset, soffset, aux=aux_attr)
-
-    return load_op.result
+    offset = _add_soffset_bytes(offset, soffset_bytes)
+    src = fx.make_view(rsrc + fx.Int32(offset), fx.make_layout(vec_width * element_bytes, 1))
+    atom = fx.make_copy_atom(fx.rocdl.BufferCopy(vec_width * dtype.width, cache_modifier), fx.Int8)
+    reg = fx.make_rmem_tensor(fx.make_layout(vec_width, 1), fx.Numeric.from_ir_type(dtype))
+    fx.copy(atom, src, reg)
+    loaded = fx.memref_load_vec(reg)
+    result = _unwrap_value(loaded[0] if vec_width == 1 else loaded)
+    if result.type != result_type:
+        result = _unwrap_value(std_arith.BitcastOp(result_type, result).result)
+    return result
 
 
 @dsl_loc_tracing
@@ -568,19 +514,20 @@ def buffer_store(
     data: ir.Value,
     rsrc: ir.Value,
     offset: ir.Value,
-    mask: Optional[ir.Value] = None,
+    mask: ir.Value | None = None,
     cache_modifier: int = 0,
     *,
-    soffset_bytes: Optional[Union[int, ir.Value]] = None,
+    soffset_bytes: int | ir.Value | None = None,
     offset_is_bytes: bool = False,
 ):
     """AMD buffer store operation.
 
     Store data to global memory using buffer descriptor and offset.
 
+
     Args:
         data: Data to store (scalar or vector)
-        rsrc: Buffer resource descriptor (!llvm.ptr<8>)
+        rsrc: Buffer descriptor pointer from :func:`create_buffer_resource`
         offset: Offset in elements (i32 type)
         mask: Optional mask for predicated store (i1 type)
         cache_modifier: Cache control flags (0 for default)
@@ -591,7 +538,8 @@ def buffer_store(
         >>> # Store with mask
         >>> buffer_store(data, rsrc, offset, mask=valid)
     """
-    # Unwrap all inputs (accept DSL Numeric values via ir_value())
+    # Unwrap all inputs (accept DSL Numeric values via ir_value()). `rsrc` stays a DSL
+    # pointer: the copy atom needs the !fly.ptr<i8, BufferDesc> fat pointer, not a value.
     if hasattr(data, "ir_value"):
         data = data.ir_value()
     if isinstance(offset, int):
@@ -599,25 +547,25 @@ def buffer_store(
     elif hasattr(offset, "ir_value"):
         offset = offset.ir_value()
     data = _unwrap_value(data)
-    rsrc = _unwrap_value(rsrc)
     offset = _unwrap_value(offset)
 
     # Convert offset to i32 if needed
-    if not isinstance(offset.type, ir.IntegerType) or offset.type.width != 32:
-        op = std_arith.IndexCastOp(T.i32(), offset)
-        offset = _unwrap_value(op.result)
-
-    # IMPORTANT: RawPtrBufferStoreOp offset is in BYTES.
+    offset = fx.Int32(offset).ir_value()
+
+    # Get element size from data type
+    data_type = data.type
+    if hasattr(data_type, "element_type"):  # Vector type
+        element_type = data_type.element_type
+        vec_width = data_type.shape[0]
+    else:  # Scalar type
+        element_type = data_type
+        vec_width = 1
+    element_bytes = element_type.width // 8
+
+    # IMPORTANT: the buffer store offset is in BYTES.
     # For backward compat, `buffer_store()` accepts element offsets by default
     # and scales them to bytes. Set `offset_is_bytes=True` to skip scaling.
     if not offset_is_bytes:
-        # Get element size from data type
-        data_type = data.type
-        if hasattr(data_type, "element_type"):  # Vector type
-            element_type = data_type.element_type
-        else:  # Scalar type
-            element_type = data_type
-        element_bytes = element_type.width // 8
         bytes_const = _create_i32_constant(element_bytes)
         op = std_arith.MulIOp(offset, bytes_const)
         offset = _unwrap_value(op.result)
@@ -629,18 +577,9 @@ def buffer_store(
         op = std_arith.SelectOp(mask, offset, max_offset)
         offset = _unwrap_value(op.result)
 
-    # Create instruction offset (soffset) and aux flags
-    if soffset_bytes is None:
-        soffset = _create_i32_constant(0)
-    else:
-        if isinstance(soffset_bytes, int):
-            soffset = _create_i32_constant(int(soffset_bytes))
-        else:
-            soffset = _unwrap_value(soffset_bytes)
-            if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32:
-                op = std_arith.IndexCastOp(T.i32(), soffset)
-                soffset = _unwrap_value(op.result)
-    aux_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) if cache_modifier else None
-
-    # Emit buffer store
-    rocdl.RawPtrBufferStoreOp(data, rsrc, offset, soffset, aux=aux_attr)
+    offset = _add_soffset_bytes(offset, soffset_bytes)
+    dst = fx.make_view(rsrc + fx.Int32(offset), fx.make_layout(vec_width * element_bytes, 1))
+    atom = fx.make_copy_atom(fx.rocdl.BufferCopy(vec_width * element_type.width, cache_modifier), fx.Int8)
+    reg = fx.make_rmem_tensor(fx.make_layout(vec_width, 1), fx.Numeric.from_ir_type(element_type))
+    fx.memref_store_vec(data, reg)
+    fx.copy(atom, reg, dst)
diff --git a/kernels/conv/conv3d_implicit.py b/kernels/conv/conv3d_implicit.py
index 454145fbcb..c64409fa54 100644
--- a/kernels/conv/conv3d_implicit.py
+++ b/kernels/conv/conv3d_implicit.py
@@ -3,8 +3,12 @@
 
 """Double-buffered implicit-GEMM conv3d (BF16).
 
-x: (N, C, D, H, W) bf16 NCDHW, weight: (K, C, T, R, S) bf16 KCTRS.
-Returns (N, K, Do, Ho, Wo) bf16. Supports stride, padding, bias, and split-K.
+x: (N, C, D, H, W) bf16 NCDHW by default, weight: (K, C/groups, T, R, S) bf16 KCTRS.
+Returns (N, K, Do, Ho, Wo) bf16 by default. ``input_layout`` / ``output_layout`` select
+NCDHW or NDHWC independently; the GEMM itself is channels-last, so NDHWC input skips the
+pre-transpose and NDHWC output is the raw row-major (npq, K) the epilogue produces.
+Supports stride, padding (int, per-axis tuple, or torch's "same" / "valid"),
+padding_mode, dilation, bias, groups, and split-K.
 """
 
 import functools
@@ -15,9 +19,8 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl._mlir.dialects import llvm
+from flydsl.compiler.kernel_function import CompilationContext
 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.mem_ops import buffer_atomic_add
 
@@ -25,6 +28,15 @@
 STAGES = 2
 WARP_SIZE = 64
 
+# K tiles consumed between two barriers. Each one is MI_M * MI_N MFMAs, and that product
+# is the only thing that hides global latency here -- see the PIPE_STAGES comment for why
+# the pipeline depth cannot. Costs no LDS (the tiles are stages that already exist) and no
+# extra ds_read/DMA traffic; it just halves the number of barriers. Worth +8..16% on the
+# 3x3 conv2d/conv3d shapes at 2. Reaching the same ratio through TILE_K = 64 instead is a
+# trap: it makes the LDS row stride 128B, exactly one bank rotation, and the resulting
+# ds_read_b128 conflicts cost more than the batching wins (measured ~15% slower).
+TILES_PER_BARRIER = 2
+
 MFMA_M = 16
 MFMA_N = 16
 MFMA_A_VALUES = 8
@@ -37,6 +49,27 @@
 
 DEFAULT_TILE = (128, 128, 2, 4)
 
+TILE_LADDER = ((128, 128, 2, 4), (64, 64, 2, 2), (32, 32, 1, 2))
+
+TILE_MIN_WAVES_PER_CU = 6
+
+PADDING_MODES = ("zeros", "reflect", "replicate", "circular")
+
+CONV_COMPILE_HINTS = {}
+
+
+def _as_stream(stream):
+    return stream if hasattr(stream, "_is_stream_param") else fx.Stream(stream)
+
+
+def _dispatch(exe, *args, stream=None):
+    """Run a builder's launcher, pre-compiling on first use."""
+    cf = getattr(exe, "_cf", None)
+    if cf is None:
+        exe._cf = exe.compile(*args, stream=stream)
+        return
+    cf(*args, _as_stream(stream))
+
 
 def _autotune_enabled():
     return os.environ.get("FLYDSL_CONV3D_AUTOTUNE", "0").lower() in ("1", "true", "yes")
@@ -49,15 +82,63 @@ def _pad_channels(c):
     return (c + LDG_VEC - 1) // LDG_VEC * LDG_VEC
 
 
+# Layout names accepted per spatial rank.
+LAYOUTS = {
+    3: ("NCDHW", "NDHWC"),
+    2: ("NCHW", "NHWC"),
+    1: ("NCW", "NWC"),
+}
+
+
+def _check_layouts(rank, input_layout, output_layout):
+    names = LAYOUTS[rank]
+    for what, v in (("input_layout", input_layout), ("output_layout", output_layout)):
+        assert v in names, f"{what} must be one of {names}, got {v!r}"
+
+
+def _shape_ncdhw(x, ndhwc):
+    """Unpack a 5-D input in either layout to (n, c, d, h, w)."""
+    if ndhwc:
+        n, d, h, w, c = x.shape
+    else:
+        n, c, d, h, w = x.shape
+    return n, c, d, h, w
+
+
+def _pad_spatial(x, ndhwc, pads, mode="constant"):
+    """Pad (D, H, W) with torch's (w_lo, w_hi, h_lo, h_hi, d_lo, d_hi) ordering."""
+    if mode == "constant":
+        return torch.nn.functional.pad(x, ((0, 0) + pads) if ndhwc else pads), ndhwc
+    if ndhwc:
+        x = x.permute(0, 4, 1, 2, 3)
+    return torch.nn.functional.pad(x, pads, mode=mode), False
+
+
+def _big_in(n, c, groups, d, h, w, pt, ph, pw):
+    """Whether the kernel's 64-bit BIG_IN address path would engage for this input."""
+    cp = _pad_channels(c // groups) * groups
+    return n * cp * (d + 2 * pt) * (h + 2 * ph) * (w + 2 * pw) > 0x7FFFFFFF
+
+
+def _evict_weight(key, _ref):
+    """weakref callback: drop the entry the dead weight was pinning."""
+    ent = _WEIGHT_CACHE.get(key)
+    if ent is not None and ent[0]() is None:
+        del _WEIGHT_CACHE[key]
+
+
 def _prep_weight(w, k, kt, kh, kw, c):
-    key = id(w)
+    """Pack (K, C, T, R, S) -> (K, T*R*S*Cpad), memoized on the source weight."""
+    anchor = w._base if w._base is not None else w
+    key = w.data_ptr()
+    stamp = (w._version, tuple(w.shape), w.stride(), w.dtype)
     ent = _WEIGHT_CACHE.get(key)
-    if ent is not None and ent[0]() is w:
+    if ent is not None and ent[0]() is anchor and ent[2] == stamp:
         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)
+    _WEIGHT_CACHE[key] = (weakref.ref(anchor, functools.partial(_evict_weight, key)), wk, stamp)
     return wk
 
 
@@ -69,6 +150,8 @@ def _prep_weight(w, k, kt, kh, kw, c):
 _TR_PAD = 8
 _TR_LDS_S = TR_TILE + _TR_PAD
 
+TR_MAX_BIG_S = (0x7FFFFFFF - (TR_TILE - TR_VEC)) // (TR_TILE - 1)
+
 
 @functools.lru_cache(maxsize=64)
 def compile_transpose_ncdhw_ndhwc(n, c, s):
@@ -104,15 +187,15 @@ class BF16Ty:
         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)
+            GPtrTy = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Global, 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)
+                return _flat_div(fx.rocdl.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_base_elem = fx.Int64(nb) * fx.Int64(c) * fx.Int64(s) + fx.Int64(c0) * fx.Int64(s) + fx.Int64(s0)
             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_base_elem = fx.Int64(nb) * fx.Int64(s) * fx.Int64(c) + fx.Int64(s0) * fx.Int64(c) + fx.Int64(c0)
             out_div = _rebased(out, out_base_elem)
         else:
             in_base = nb * c * s
@@ -122,9 +205,11 @@ def _rebased(tensor, base_elem):
         tr_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_ty)
         tr_reg = fx.make_rmem_tensor(TR_VEC, elem_ty)
 
+        _lds_st_ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Shared, TR_VEC * BF16_BYTES)
+
         def lds_store_vec8(elem_offset, value):
-            ptr = (fx.recast_iter(fx.Uint8, lds.ptr) + fx.Int32(elem_offset * 2)).llvm_ptr
-            llvm.StoreOp(value, ptr, alignment=16)
+            base = fx.Int64(fx.ptrtoint(lds.ptr)) + fx.Int64(elem_offset * 2)
+            fx.ptr_store(value, fx.inttoptr(_lds_st_ptr_ty, base))
 
         def lds_load_scalar(elem_offset):
             u8 = fx.recast_iter(fx.Uint8, lds.ptr)
@@ -145,9 +230,10 @@ def lds_load_scalar(elem_offset):
             safe = arith.select(valid, g, fx.Int32(0))
             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)
+            lds_store_vec8(rc * _TR_LDS_S + sv, v)
 
-        llvm.InlineAsmOp(None, [], "s_waitcnt lgkmcnt(0)\n\ts_barrier", "", has_side_effects=True)
+        rocdl.s_waitcnt(lgkmcnt=0)
+        rocdl.s_barrier()
 
         for i in range_constexpr(_TR_ITERS):
             lin = tid + i * TR_THREADS
@@ -174,24 +260,60 @@ def launch_transpose(out: fx.Tensor, inp: fx.Tensor, stream: fx.Stream = fx.Stre
             stream=stream,
         )
 
-    return launch_transpose
+    def _launch(out, inp, stream=None):
+        with CompilationContext.compile_hints(CONV_COMPILE_HINTS):
+            return launch_transpose(out, inp, stream=_as_stream(stream))
+
+    def _compile(out, inp, stream=None):
+        with CompilationContext.compile_hints(CONV_COMPILE_HINTS):
+            return flyc.compile(launch_transpose, out, inp, _as_stream(stream))
+
+    _launch.compile = _compile
+    return _launch
 
 
 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
+    big = n * c * s > 0x7FFFFFFF
     if not (x.is_contiguous() and x.dtype == torch.bfloat16 and c % TR_VEC == 0):
         return x.permute(0, 2, 3, 4, 1).contiguous()
+    if big and s > TR_MAX_BIG_S:
+        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)
+    _dispatch(exe, out, x, stream=torch.cuda.current_stream() if stream is None else stream)
     return out
 
 
 @functools.lru_cache(maxsize=256)
 def compile_conv3d_implicit(
-    n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias=False, splitk=1, tile=DEFAULT_TILE, wgm=1
+    n,
+    c,
+    d,
+    h,
+    w,
+    k,
+    kt,
+    kh,
+    kw,
+    st,
+    sh,
+    sw,
+    pt,
+    ph,
+    pw,
+    dt=1,
+    dh=1,
+    dw=1,
+    pad_mode="zeros",
+    has_bias=False,
+    splitk=1,
+    tile=DEFAULT_TILE,
+    wgm=1,
+    groups=1,
+    out_ndhwc=False,
 ):
     TILE_M, TILE_N, WAVE_M, WAVE_N = tile
     BLOCK_THREADS = WAVE_M * WAVE_N * WARP_SIZE
@@ -205,29 +327,38 @@ def compile_conv3d_implicit(
     LDG_A_COUNT = TILE_M * TILE_K // BLOCK_VECS
     LDG_B_COUNT = TILE_N * TILE_K // BLOCK_VECS
 
+    # `c` is the padded TOTAL channel count and stays the NDHWC row stride. CGP is the
+    # per-group channel count and is what the GEMM K axis decomposes against; the two
+    # coincide only when groups == 1.
+    CGP = c // groups
+    KG = k // groups
+
     assert TILE_K == 32
     assert TILE_M % (WAVE_M * MFMA_M) == 0, f"TILE_M={TILE_M} not divisible by WAVE_M*16"
     assert TILE_N % (WAVE_N * MFMA_N) == 0, f"TILE_N={TILE_N} not divisible by WAVE_N*16"
     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, f"c={c} must be a multiple of LDG_VEC={LDG_VEC}; use _conv3d_impl to pad"
+    assert c % groups == 0, f"c={c} not divisible by groups={groups}"
+    assert k % groups == 0, f"k={k} not divisible by groups={groups}"
+    assert CGP % LDG_VEC == 0, f"c/groups={CGP} 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
-    ho = (h + 2 * ph - kh) // sh + 1
-    wo = (w + 2 * pw - kw) // sw + 1
+    # Dilation only stretches the filter's footprint; the K axis (CRS) is unchanged.
+    do = (d + 2 * pt - (dt * (kt - 1) + 1)) // st + 1
+    ho = (h + 2 * ph - (dh * (kh - 1) + 1)) // sh + 1
+    wo = (w + 2 * pw - (dw * (kw - 1) + 1)) // sw + 1
     dhw = do * ho * wo
     hw_o = ho * wo
     npq = n * dhw
-    crs = c * kt * kh * kw
+    crs = CGP * kt * kh * kw
     k_tiles = (crs + TILE_K - 1) // TILE_K
 
     BIG_IN = (n * c * d * h * w) > 0x7FFFFFFF
     BIG_OUT = (n * k * do * ho * wo * BF16_BYTES) > 0x7FFFFFFF
 
     X_BYTES = n * c * d * h * w * BF16_BYTES
-    W_BYTES = k * c * kt * kh * kw * BF16_BYTES
+    W_BYTES = k * crs * BF16_BYTES
     OOB_SENTINEL_ELEM = 0x7FFFFF80  # *2 = 0xFFFFFF00 bytes (~4.2950 GB), just under 2^32
     OOB_SENTINEL_BYTES = OOB_SENTINEL_ELEM * BF16_BYTES
     BIG_IN_NR = 0x80000000  # 2 GB num_records for the rebased BIG_IN resource
@@ -235,28 +366,60 @@ def compile_conv3d_implicit(
     assert X_BYTES < OOB_SENTINEL_BYTES or BIG_IN, f"input {X_BYTES}B exceeds limit"
     BIG_IN_N1 = BIG_IN and n == 1
     BIG_IN_NM = BIG_IN and n > 1
-    X_SAMPLE_BYTES = c * d * h * w * BF16_BYTES
 
-    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
+    _t_aligned = BIG_IN_N1 and hw_o % TILE_M == 0
+    if BIG_IN_N1:
+        _ot_span = (TILE_M - 1) // hw_o + (1 if _t_aligned else 2)
+        _t_span = min(d - 1, (_ot_span - 1) * st + dt * (kt - 1))
+        _h_span = min(h - 1, ((TILE_M - 1) // wo + 1) * sh + dh * (kh - 1)) if _t_aligned else h - 1
+        _span = (((_t_span * h + _h_span) * w + (w - 1)) * c + c) * BF16_BYTES
+        assert _span <= BIG_IN_NR, (
+            f"input sample too large for the 32-bit gather: a {TILE_M}-row tile reaches "
+            f"{_span / 2**30:.2f} GiB from its rebased origin, past the "
+            f"{BIG_IN_NR / 2**30:.0f} GiB the buffer descriptor addresses. Split the batch "
+            f"over N, or pass a narrower tile=(TILE_M, ...)."
+        )
+
+    assert pad_mode in PADDING_MODES, f"pad_mode must be one of {PADDING_MODES}, got {pad_mode!r}"
+    assert pad_mode == "zeros" or not BIG_IN, "non-zero pad_mode requires the non-BIG_IN address path"
+    X_SAMPLE_ELEMS = c * d * h * w
+
+    tiles_per_group = (KG + TILE_N - 1) // TILE_N
+    n_tail = KG % TILE_N != 0
+    grid_n = groups * tiles_per_group
 
     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
-    # hide DMA latency; a shallower pipeline measured slower (2/3/4-stage A/B).
-    PIPE_STAGES = 4
+    Y_BYTES = npq * k * (4 if use_splitk else BF16_BYTES)
+
+    assert (
+        not use_splitk or npq * k * 4 <= SPLITK_MAX_STAGING_BYTES
+    ), f"split-K staging {npq * k * 4}B exceeds the {SPLITK_MAX_STAGING_BYTES}B buffer window"
+
+    PIPE_STAGES = 2 * TILES_PER_BARRIER
 
     LDS_A_SIZE = PIPE_STAGES * TILE_M * TILE_K
     LDS_B_SIZE = PIPE_STAGES * TILE_N * TILE_K
 
     grid_m = (npq + TILE_M - 1) // TILE_M
-    WGM = max(1, int(wgm))
+
+    MAX_GRID_X = 0xFFFFFFFF // BLOCK_THREADS
+    MAX_GRID_YZ = 65535
+    grid_x = min(grid_m, MAX_GRID_X)
+    m_chunks = (grid_m + grid_x - 1) // grid_x
+
+    _row_chk = (npq % TILE_M != 0) or (grid_x * m_chunks > grid_m)
+    _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) and (not out_ndhwc)
+
+    assert grid_n <= MAX_GRID_YZ, f"grid.y = {grid_n} exceeds the {MAX_GRID_YZ}-block limit"
+    assert (
+        m_chunks * splitk <= MAX_GRID_YZ
+    ), f"grid.z = {m_chunks} M-chunks x {splitk} splits exceeds the {MAX_GRID_YZ}-block limit"
+
+    WGM = 1 if m_chunks > 1 else max(1, int(wgm))
     elem_ty = fx.BFloat16
     mfma_fn = rocdl.mfma_f32_16x16x32_bf16
     temporal_only_fast = (
@@ -274,7 +437,26 @@ 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_buf = fx.rocdl.make_buffer_tensor(y)
+        # The im2col gather addresses its source by a flat element index, so the DMA source
+        # is a 1-D view over the buffer rather than the tensor's own n-D layout: dividing
+        # that by a 1-element tile makes slice(src, (None, off)) exactly element `off`,
+        # with no coordinate decomposition. `elems` only shapes the view -- the sentinel
+        # offset that masks a tap deliberately points past it, and num_records (not the
+        # layout) is what turns that into a zero-fill.
+        def _dma_src(ptr, elems, num_records_bytes):
+            buf = fx.rocdl.make_buffer_ptr(ptr, num_records_bytes=num_records_bytes)
+            return fx.logical_divide(fx.make_view(buf, fx.make_layout(elems, 1)), fx.make_layout(1, 1))
+
+        def _x_rebased(off_elems):
+            # BIG_IN moves the descriptor base to the block's own origin so a 32-bit
+            # voffset still reaches the tile; num_records bounds it at 2 GB from there.
+            ptr = fx.add_offset(fx.get_iter(x), fx.make_int_tuple(off_elems))
+            return _dma_src(ptr, BIG_IN_NR // BF16_BYTES, BIG_IN_NR)
+
+        w_src = _dma_src(fx.get_iter(weight), W_BYTES // BF16_BYTES, W_BYTES)
+        if const_expr(not BIG_IN):
+            x_src = _dma_src(fx.get_iter(x), X_BYTES // BF16_BYTES, X_BYTES)
+        y_buf = fx.rocdl.make_buffer_tensor(y, num_records_bytes=Y_BYTES)
         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))
@@ -288,15 +470,6 @@ def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias:
             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)
-        w_buf = fx.Tensor(fx.make_view(fx.get_iter(w_buf0), fx.make_layout(k * crs, 1)))
-        w_div = fx.logical_divide(w_buf, fx.make_layout(1, 1))
-        if const_expr(not BIG_IN):
-            x_buf0 = fx.rocdl.make_buffer_tensor(x, max_size=False)
-            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_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)
@@ -307,46 +480,56 @@ def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias:
         b_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_B_SIZE, 16]).peek()
 
         tid = fx.thread_idx.x
-        if const_expr(WGM > 1):
-            pid = fx.Index(fx.block_idx.x) + fx.Index(fx.block_idx.y) * fx.Index(grid_m)
-            blocks_per_group = fx.Index(WGM * grid_n)
-            group_id = pid // blocks_per_group
-            first_m = group_id * fx.Index(WGM)
-            group_rows = fx.Index(grid_m) - first_m
-            group_rows = fx.Index(arith.select(group_rows < fx.Index(WGM), group_rows, fx.Index(WGM)))
-            local = pid % blocks_per_group
-            m_offset = fx.Index(first_m + (local % group_rows)) * TILE_M
-            n_offset = fx.Index(local // group_rows) * TILE_N
+        if const_expr(m_chunks > 1):
+            m_chunk = fx.Int64(fx.block_idx.z) % fx.Int64(m_chunks)
+            m_offset = (fx.Int64(fx.block_idx.x) + m_chunk * fx.Int64(grid_x)) * TILE_M
+            n_tile = fx.block_idx.y
+        elif const_expr(WGM > 1):
+            pid = fx.Int64(fx.block_idx.x) + fx.Int64(fx.block_idx.y) * fx.Int64(grid_m)
+            blocks_per_swizzle = fx.Int64(WGM * grid_n)
+            swizzle_id = pid // blocks_per_swizzle
+            first_m = swizzle_id * fx.Int64(WGM)
+            swizzle_rows = fx.Int64(grid_m) - first_m
+            swizzle_rows = fx.Int64(arith.select(swizzle_rows < fx.Int64(WGM), swizzle_rows, fx.Int64(WGM)))
+            local = pid % blocks_per_swizzle
+            m_offset = fx.Int64(first_m + (local % swizzle_rows)) * TILE_M
+            n_tile = fx.Int64(local // swizzle_rows)
         else:
             m_offset = fx.block_idx.x * TILE_M
-            n_offset = fx.block_idx.y * TILE_N
+            n_tile = fx.block_idx.y
+
+        if const_expr(groups > 1):
+            gi = n_tile // tiles_per_group
+            n_local = (n_tile % tiles_per_group) * TILE_N
+            n_offset = gi * KG + n_local
+            ch_base = gi * CGP
+        else:
+            n_offset = n_tile * TILE_N
+            n_local = n_offset
         if const_expr(use_splitk):
-            k_off = fx.block_idx.z * (tiles_per_split * TILE_K)
+            if const_expr(m_chunks > 1):
+                split_idx = fx.Int64(fx.block_idx.z) // fx.Int64(m_chunks)
+            else:
+                split_idx = fx.Int64(fx.block_idx.z)
+            k_off = split_idx * (tiles_per_split * TILE_K)
         else:
             k_off = 0
 
-        # 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)
-            buf_ptr = make_buffer_ptr(gptr, num_records_bytes=BIG_IN_NR)
-            # 1-D element view so the flat im2col gather offset indexes elements.
-            buf = fx.Tensor(fx.make_view(buf_ptr, fx.make_layout(BIG_IN_NR // BF16_BYTES, 1)))
-            return fx.logical_divide(buf, fx.make_layout(1, 1))
-
         if const_expr(BIG_IN_N1):
             nbase = m_offset // dhw
-            ot_base0 = (m_offset % dhw) // hw_o
-            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(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(fx.ptrtoint(fx.get_iter(x)))
+            rem0 = m_offset % dhw
+            ot_base0 = rem0 // hw_o
+
+            base_t = ot_base0 * fx.Int64(st) - fx.Int64(pt)
+            base_t = arith.select(base_t < fx.Int64(0), fx.Int64(0), base_t)
+            if const_expr(_t_aligned):
+                oh_base0 = (rem0 % hw_o) // wo
+                base_h = oh_base0 * fx.Int64(sh) - fx.Int64(ph)
+                base_h = arith.select(base_h < fx.Int64(0), fx.Int64(0), base_h)
+            else:
+                base_h = fx.Int64(0)
+            x_base_elem = ((nbase * fx.Int64(d) + base_t) * fx.Int64(h) + base_h) * fx.Int64(w) * fx.Int64(c)
+            x_src = _x_rebased(fx.Int64(x_base_elem))
 
         wid = tid // WARP_SIZE
         lane = tid % WARP_SIZE
@@ -369,26 +552,56 @@ class Vec8Ty:
         acc = [acc0 for _ in range_constexpr(N_ACC)]
 
         def barrier(vmcnt=0, lgkmcnt=None):
-            waits = []
-            if vmcnt is not None:
-                waits.append(f"vmcnt({vmcnt})")
-            if lgkmcnt is not None:
-                waits.append(f"lgkmcnt({lgkmcnt})")
-            pre = ("s_waitcnt " + " ".join(waits) + "\n\t") if waits else ""
-            llvm.InlineAsmOp(None, [], f"{pre}s_barrier", "", has_side_effects=True)
+            rocdl.s_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)
+            rocdl.s_barrier()
 
         def lds_load_vec8(lds_array, elem_offset):
             u8_ptr = fx.recast_iter(fx.Uint8, lds_array.ptr)
             return fx.ptr_load(u8_ptr + fx.Int32(elem_offset * 2), result_type=Vec8Ty)
 
         def a_lds_off(stage, row, col):
-            return (fx.Index(stage) * TILE_M + row) * TILE_K + col
+            return (fx.Int64(stage) * TILE_M + row) * TILE_K + col
 
         def b_lds_off(stage, row, col):
-            return (fx.Index(stage) * TILE_N + row) * TILE_K + col
+            return (fx.Int64(stage) * TILE_N + row) * TILE_K + col
 
         def in_range(v, hi):
-            return (v >= 0) & (v < fx.Index(hi))
+            return (v >= 0) & (v < fx.Int64(hi))
+
+        def dil(tap, factor):
+            scaled = tap * factor if const_expr(factor != 1) else tap
+            return scaled
+
+        def pad_coord(v, ext, pad):
+            """Tap coordinate -> in-bounds input coordinate; returns (coord, mask).
+
+            "zeros" leaves the coordinate alone and returns a range mask, which the
+            caller folds into the OOB-sentinel routing so the load reads as zero. Every
+            other mode resolves the coordinate into [0, ext) instead and returns no mask
+            """
+            if const_expr(pad_mode == "zeros"):
+                return v, in_range(v, ext)
+            u = v + fx.Int64(pad)
+            low = u < fx.Int64(pad)  # v < 0
+            high = u >= fx.Int64(pad + ext)  # v >= ext
+            mid = u - fx.Int64(pad)  # v, where in range
+            if const_expr(pad_mode == "replicate"):
+                r = arith.select(high, fx.Int64(ext - 1), mid)
+                r = arith.select(low, fx.Int64(0), r)
+            elif const_expr(pad_mode == "reflect"):
+                # [a b c d e] pad 2 -> [c b a b c d e d c]: -v near, 2*(ext-1) - v far.
+                r = arith.select(high, fx.Int64(2 * (ext - 1) + pad) - u, mid)
+                r = arith.select(low, fx.Int64(pad) - u, r)
+            else:  # circular: v + ext near, v - ext far
+                r = arith.select(high, u - fx.Int64(pad + ext), mid)
+                r = arith.select(low, u + fx.Int64(ext - pad), r)
+            return fx.Int64(r), None
+
+        def gather_valid(base, *masks):
+            for m in masks:
+                if const_expr(m is not None):
+                    base = base & m
+            return base
 
         # ---- Per-thread row decomposition (loop-invariant across K) ----
         _row_dec = []  # per-i tuple of precomputed row terms
@@ -397,7 +610,7 @@ def in_range(v, hi):
             local_m = linear // TILE_K
             local_k = linear % TILE_K
             row = m_offset + local_m
-            row_valid = row < fx.Index(npq)
+            row_valid = row < fx.Int64(npq)
             if const_expr(temporal_only_fast):
                 out_t = (row // hw_o) % d
                 _row_dec.append((local_k, row, row_valid, out_t))
@@ -419,55 +632,62 @@ def in_range(v, hi):
                 else:
                     _row_dec.append((local_k, row_valid, n_idx, in_t0, in_h0, in_w0))
 
-        SCALAR_K = c % TILE_K == 0
+        SCALAR_K = CGP % TILE_K == 0
 
         # ---- 3D im2col address math ----
+        # The K axis decomposes against CGP (per-group channels) while every g_off below
+        # keeps `c` (padded total channels) as the NDHWC row stride. `cc` is the absolute
+        # input channel: the group base plus the offset within the group.
         def _a_addr(i, kbase_i, cc_base, ckk_base):
             dec = _row_dec[i]
             local_k = dec[0]
-            k_abs = kbase_i + fx.Index(local_k)
+            k_abs = kbase_i + fx.Int64(local_k)
             if const_expr(SCALAR_K):
-                cc = cc_base + fx.Index(local_k)
+                cc = cc_base + fx.Int64(local_k)  # cc_base already carries ch_base
             else:
-                cc = k_abs % c
-            k_valid = k_abs < fx.Index(crs)
+                cc = k_abs % CGP
+                if const_expr(groups > 1):
+                    cc = ch_base + cc
+            k_valid = k_abs < fx.Int64(crs)
             if const_expr(temporal_only_fast):
                 _, row, row_valid, out_t = dec
-                kt_i = ckk_base if const_expr(SCALAR_K) else k_abs // c
-                temporal_delta = kt_i - pt
-                in_t = out_t + temporal_delta
-                valid = row_valid & k_valid & in_range(in_t, d)
+                kt_i = ckk_base if const_expr(SCALAR_K) else k_abs // CGP
+                temporal_delta = dil(kt_i, dt) - pt
+                in_t, m_t = pad_coord(out_t + temporal_delta, d, pt)
+                valid = gather_valid(row_valid & k_valid, m_t)
+
+                delta = temporal_delta if const_expr(pad_mode == "zeros") else (in_t - out_t)
                 if const_expr(BIG_IN_N1):
-                    g_off = ((row + temporal_delta * hw_o) - (fx.Index(nbase) * dhw + base_t * hw_o)) * c + cc
+                    g_off = ((row + delta * hw_o) - (fx.Int64(nbase) * dhw + base_t * hw_o)) * c + cc
                 else:
-                    g_off = (row + temporal_delta * hw_o) * c + cc
+                    g_off = (row + delta * hw_o) * c + cc
             else:
-                ckk = ckk_base if const_expr(SCALAR_K) else k_abs // c
+                ckk = ckk_base if const_expr(SCALAR_K) else k_abs // CGP
                 kw_i = ckk % kw
                 ckk2 = ckk // kw
                 kh_i = ckk2 % kh
                 kt_i = ckk2 // kh
                 if const_expr(BIG_IN_N1):
                     _, row_valid, di, in_t0, in_h0, in_w0 = dec
-                    in_t = in_t0 + kt_i
-                    in_h = in_h0 + kh_i
-                    in_w = in_w0 + kw_i
-                    valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w)
-                    g_off = (((di * d + (in_t - base_t)) * h + in_h) * w + in_w) * c + cc
+                    in_t, m_t = pad_coord(in_t0 + dil(kt_i, dt), d, pt)
+                    in_h, m_h = pad_coord(in_h0 + dil(kh_i, dh), h, ph)
+                    in_w, m_w = pad_coord(in_w0 + dil(kw_i, dw), w, pw)
+                    valid = gather_valid(row_valid & k_valid, m_t, m_h, m_w)
+                    g_off = (((di * d + (in_t - base_t)) * h + (in_h - base_h)) * w + in_w) * c + cc
                 elif const_expr(BIG_IN_NM):
                     _, row_valid, n_idx, in_t0, in_h0, in_w0 = dec
-                    in_t = in_t0 + kt_i
-                    in_h = in_h0 + kh_i
-                    in_w = in_w0 + kw_i
-                    valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w)
+                    in_t, m_t = pad_coord(in_t0 + dil(kt_i, dt), d, pt)
+                    in_h, m_h = pad_coord(in_h0 + dil(kh_i, dh), h, ph)
+                    in_w, m_w = pad_coord(in_w0 + dil(kw_i, dw), w, pw)
+                    valid = gather_valid(row_valid & k_valid, m_t, m_h, m_w)
                     g_off = ((in_t * h + in_h) * w + in_w) * c + cc
                     return fx.Int32(g_off), valid, n_idx
                 else:
                     _, row_valid, n_idx, in_t0, in_h0, in_w0 = dec
-                    in_t = in_t0 + kt_i
-                    in_h = in_h0 + kh_i
-                    in_w = in_w0 + kw_i
-                    valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w)
+                    in_t, m_t = pad_coord(in_t0 + dil(kt_i, dt), d, pt)
+                    in_h, m_h = pad_coord(in_h0 + dil(kh_i, dh), h, ph)
+                    in_w, m_w = pad_coord(in_w0 + dil(kw_i, dw), w, pw)
+                    valid = gather_valid(row_valid & k_valid, m_t, m_h, m_w)
                     g_off = (((n_idx * d + in_t) * h + in_h) * w + in_w) * c + cc
             return fx.Int32(g_off), valid
 
@@ -475,83 +695,80 @@ def _b_addr(i, k_base):
             linear = (tid + i * BLOCK_THREADS) * LDG_VEC
             local_n = linear // TILE_K
             local_k = linear % TILE_K
-            col = n_offset + fx.Index(local_n)
-            g_off = fx.Int32(col * crs + (fx.Index(k_base) + fx.Index(local_k)))
-            col_valid = (col < fx.Index(k)) if const_expr(n_tail) else None
+            col = n_offset + fx.Int64(local_n)
+            g_off = fx.Int32(col * crs + (fx.Int64(k_base) + fx.Int64(local_k)))
+            # Tail is per group: the N grid is over-provisioned to groups*tiles_per_group.
+            col_valid = ((n_local + fx.Int64(local_n)) < fx.Int64(KG)) if const_expr(n_tail) else None
             return g_off, col_valid
 
         # ---- global -> LDS DMA copy, masking via OOB routing ----
-        # OOB sentinel element indices: a gather past num_records makes the buffer
-        # hardware return 0, reproducing the padding/halo zeroing.
-        X_ELEMS = fx.Int32(n * c * d * h * w)
-        W_ELEMS = fx.Int32(k * c * kt * kh * kw)
-        BIG_OOB_ELEM = fx.Int32(BIG_IN_NR // BF16_BYTES)
-
-        # global->LDS DMA (BufferCopyLDS128b): gather offset as a flat element index;
-        # OOB padding routes past-end. Used for both non-BIG and BIG_IN gathers.
-        g2s_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128)
-        LdsPtrTy = fx.PointerType.get(elem_ty.ir_type, 2, 512)
-
-        def _copy_g2s(src_div, lds_array, stage_tile, i, src_elem):
-            off_elems = fx.Int32(stage_tile) + (fx.Int32(tid) + fx.Int32(i * BLOCK_THREADS)) * fx.Int32(LDG_VEC)
-            lds_byte_addr = fx.Int32(fx.ptrtoint(lds_array.ptr)) + off_elems * fx.Int32(BF16_BYTES)
-            lds_ptr = fx.inttoptr(LdsPtrTy, lds_byte_addr)
-            dst = fx.make_view(lds_ptr, fx.make_layout(1, 1))
-            src = fx.slice(src_div, (None, fx.Int32(src_elem)))
-            fx.copy(g2s_atom, src, dst)
+        DMA_BYTES = LDG_VEC * BF16_BYTES  # 16
+        OOB_ELEM = fx.Int32(OOB_SENTINEL_ELEM)
+
+        _lds_dma_ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Shared, DMA_BYTES)
+
+        def sgpr(x):
+            # Hoist a wave-uniform value into an SGPR (readfirstlane).
+            return fx.Int64(rocdl.readfirstlane(T.i64, fx.Int64(x)))
+
+        _dma_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), DMA_BYTES * 8)
+
+        def _lds_dma_ptr(lds_array, stage_tile, i):
+            # buffer_load_lds takes one wave-uniform LDS base and fans the wave's lanes
+            # out from it, so the lane-0 address is the base the whole wave writes from.
+            off_elems = fx.Int64(stage_tile) + (fx.Int64(tid) + fx.Int64(i * BLOCK_THREADS)) * fx.Int64(LDG_VEC)
+            base_bytes = off_elems * fx.Int64(BF16_BYTES)
+            addr = fx.Int64(fx.ptrtoint(lds_array.ptr)) + fx.Int64(base_bytes)
+            return fx.make_view(fx.inttoptr(_lds_dma_ptr_ty, sgpr(addr)), fx.make_layout(1, 1))
+
+        def _dma_to_lds(src, dst, voff_elem):
+            fx.copy(_dma_atom, fx.slice(src, (None, voff_elem)), dst)
 
         def _load_a(stage, k_base):
-            kbase_i = fx.Index(k_base)
+            kbase_i = fx.Int64(k_base)
             cc_base = ckk_base = None
             if const_expr(SCALAR_K):
-                cc_base = kbase_i % c
-                ckk_base = kbase_i // c
-            stage_tile = fx.Index(stage) * TILE_M * TILE_K
+                cc_base = kbase_i % CGP
+                if const_expr(groups > 1):
+                    cc_base = ch_base + cc_base
+                ckk_base = kbase_i // CGP
+            stage_tile = fx.Int64(stage) * TILE_M * TILE_K
             for i in range_constexpr(LDG_A_COUNT):
                 if const_expr(BIG_IN_NM):
-                    # Rebase the buffer tensor per load to the sample base.
                     addr_ret = _a_addr(i, kbase_i, cc_base, ckk_base)
                     g_off_i, valid, n_idx_i = addr_ret
-                    sample_addr = x_base_addr + fx.Int64(n_idx_i) * fx.Int64(X_SAMPLE_BYTES)
-                    x_div_i = _x_div_from_addr(sample_addr)
-                    voff = fx.Int32(arith.select(valid, g_off_i, BIG_OOB_ELEM))
-                    _copy_g2s(x_div_i, a_lds, stage_tile, i, voff)
-                elif const_expr(BIG_IN):
-                    # BIG_IN_N1: rebased buffer tensor (built once above).
-                    g_off_i, valid = _a_addr(i, kbase_i, cc_base, ckk_base)
-                    voff = fx.Int32(arith.select(valid, g_off_i, BIG_OOB_ELEM))
-                    _copy_g2s(x_div_big, a_lds, stage_tile, i, voff)
+                    x_src_i = _x_rebased(fx.Int64(n_idx_i) * fx.Int64(X_SAMPLE_ELEMS))
+                    voff = fx.Int32(arith.select(valid, g_off_i, OOB_ELEM))
+                    _dma_to_lds(x_src_i, _lds_dma_ptr(a_lds, stage_tile, i), voff)
                 else:
                     g_off_i, valid = _a_addr(i, kbase_i, cc_base, ckk_base)
-                    voff = fx.Int32(arith.select(valid, g_off_i, X_ELEMS))
-                    _copy_g2s(x_div, a_lds, stage_tile, i, voff)
+                    voff = fx.Int32(arith.select(valid, g_off_i, OOB_ELEM))
+                    _dma_to_lds(x_src, _lds_dma_ptr(a_lds, stage_tile, i), voff)
 
         def _load_b(stage, k_base):
-            stage_tile = fx.Index(stage) * TILE_N * TILE_K
+            stage_tile = fx.Int64(stage) * TILE_N * TILE_K
             for i in range_constexpr(LDG_B_COUNT):
                 g_off, col_valid = _b_addr(i, k_base)
                 if const_expr(n_tail):
-                    voff = fx.Int32(arith.select(col_valid, g_off, W_ELEMS))
+                    voff = fx.Int32(arith.select(col_valid, g_off, OOB_ELEM))
                 else:
                     voff = g_off
-                _copy_g2s(w_div, b_lds, stage_tile, i, voff)
+                _dma_to_lds(w_src, _lds_dma_ptr(b_lds, stage_tile, i), voff)
 
         # ---- single-vec ds_read (LDS -> register), indexed by per-wave MFMA row ----
         def read_a_vec(stage, mi):
             a_row = wave_m * WARP_M + mi * MFMA_M + lane_m
-            return lds_load_vec8(a_lds, a_lds_off(stage, fx.Index(a_row), fx.Index(lane_k_a)))
+            return lds_load_vec8(a_lds, a_lds_off(stage, fx.Int64(a_row), fx.Int64(lane_k_a)))
 
         def read_b_vec(stage, ni):
             b_row = wave_n * WARP_N + ni * MFMA_N + lane_n
-            return lds_load_vec8(b_lds, b_lds_off(stage, fx.Index(b_row), fx.Index(lane_k_b)))
+            return lds_load_vec8(b_lds, b_lds_off(stage, fx.Int64(b_row), fx.Int64(lane_k_b)))
 
         def mfma_one(a_frag, b_frag, c_frag):
-            out = mfma_fn(
+            return mfma_fn(
                 T.vec(MFMA_C_VALUES, T.f32),
                 [a_frag, b_frag, c_frag, 0, 0, 0],
             )
-            rocdl.sched_mfma(1)
-            return out
 
         def read_a_frags(stage):
             frags = [read_a_vec(stage, mi) for mi in range_constexpr(MI_M)]
@@ -569,64 +786,89 @@ def do_compute(acc_values, a_frag_values, b_frag_values):
                 for ni in range_constexpr(MI_N):
                     idx = mi * MI_N + ni
                     acc_values[idx] = mfma_one(a_frag_values[mi], b_frag_values[ni], acc_values[idx])
+                rocdl.sched_mfma(MI_N)
             rocdl.s_setprio(0)
             return acc_values
 
         # global->LDS software pipeline
         # ---- prologue: fill the pipeline with the first PREFETCH tiles' DMAs ----
-        PREFETCH = PIPE_STAGES - 1
+        PREFETCH = TILES_PER_BARRIER
         for s in range_constexpr(PREFETCH):
             if const_expr(s < tiles_per_split):
                 _load_a(s, k_off + s * TILE_K)
                 _load_b(s, k_off + s * TILE_K)
-        LDG_PER_TILE = LDG_A_COUNT + LDG_B_COUNT
-
-        # ---- main loop: wait oldest tile, read frags, launch tile PREFETCH ahead, compute ----
-        for kt_idx in range_constexpr(tiles_per_split):
-            cur = kt_idx % PIPE_STAGES
-            inflight_tiles = min(PREFETCH - 1, tiles_per_split - 1 - kt_idx)
-            barrier(vmcnt=inflight_tiles * LDG_PER_TILE, lgkmcnt=0)
-            a_frags = read_a_frags(cur)
-            b_frags = read_b_frags(cur)
-            nxt = kt_idx + PREFETCH
-            if const_expr(nxt < tiles_per_split):
-                _load_a(nxt % PIPE_STAGES, k_off + nxt * TILE_K)
-                _load_b(nxt % PIPE_STAGES, k_off + nxt * TILE_K)
-                rocdl.sched_vmem(LDG_A_COUNT + LDG_B_COUNT)
-            acc = do_compute(acc, a_frags, b_frags)
+
+        # ---- main loop
+        for kt_idx in range_constexpr(0, tiles_per_split, TILES_PER_BARRIER):
+            batch = range_constexpr(kt_idx, min(kt_idx + TILES_PER_BARRIER, tiles_per_split))
+
+            barrier(vmcnt=0, lgkmcnt=0)
+            a_frags = [read_a_frags(kt % PIPE_STAGES) for kt in batch]
+            b_frags = [read_b_frags(kt % PIPE_STAGES) for kt in batch]
+            issued = 0
+            for kt in batch:
+                nxt = kt + PREFETCH
+                if const_expr(nxt < tiles_per_split):
+                    _load_a(nxt % PIPE_STAGES, k_off + nxt * TILE_K)
+                    _load_b(nxt % PIPE_STAGES, k_off + nxt * TILE_K)
+                    issued += LDG_A_COUNT + LDG_B_COUNT
+            if const_expr(issued):
+                rocdl.sched_vmem(issued)
+            for j in range_constexpr(len(batch)):
+                acc = do_compute(acc, a_frags[j], b_frags[j])
 
         if const_expr(BIG_OUT):
             y_elem_base = fx.Int64(fx.ptrtoint(fx.get_iter(y)))
 
+        _big_st_ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Global, BF16_BYTES)
+
         def _big_store(off_nk_i64, value):
             addr = y_elem_base + off_nk_i64 * fx.Int64(BF16_BYTES)
-            ptr = fx.inttoptr(GYPtrTy, addr).llvm_ptr
-            llvm.StoreOp(value.ir_value() if hasattr(value, "ir_value") else value, ptr, alignment=2)
+            fx.ptr_store(value, fx.inttoptr(_big_st_ptr_ty, addr))
 
-        def _valid_raw(row, col):
+        def _valid_raw(row, col_loc):
             if const_expr(_row_chk and n_tail):
-                return fx.as_ir_value((row < fx.Index(npq)) & (col < fx.Index(k)))
+                return arith.andi(row < fx.Int64(npq), col_loc < fx.Int64(KG))
             if const_expr(_row_chk):
-                v = row < fx.Index(npq)
-                return fx.as_ir_value(v & v)
-            v = col < fx.Index(k)
-            return fx.as_ir_value(v & v)
+                v = row < fx.Int64(npq)
+                return arith.andi(v, v)
+            v = col_loc < fx.Int64(KG)
+            return arith.andi(v, v)
+
+        _route_store = _need_chk and not use_splitk and not BIG_OUT
+
+        def _route(off, row, col_loc):
+            if const_expr(not _route_store):
+                return fx.Int32(off)
+            return fx.Int32(arith.select(_valid_raw(row, col_loc), fx.Int32(off), OOB_ELEM))
+
+        def _cols(ni):
+            """Global out-channel for MFMA column block ni, and its index within the group."""
+            col_off = fx.Int64(wave_n * WARP_N + ni * MFMA_N + c_n)
+            col = n_offset + col_off
+            return col, ((n_local + col_off) if const_expr(groups > 1) else col)
 
         def store_acc():
+            if const_expr(has_bias and not use_splitk):
+                bias_vals = []
+                for ni in range_constexpr(MI_N):
+                    col, col_loc = _cols(ni)
+                    col_i = fx.Int32(col)  # bias is indexed by the global out-channel
+                    if const_expr(n_tail):
+                        col_i = arith.select(col_loc < fx.Int64(KG), col_i, fx.Int32(0))
+                    fx.copy(bias_atom, fx.slice(bias_div, (None, col_i)), bias_reg)
+                    bias_vals.append(fx.Float32(fx.memref_load_vec(bias_reg)[0]))
+
             for mi in range_constexpr(MI_M):
                 row_base = m_offset + wave_m * WARP_M + mi * MFMA_M + c_m_vec
                 for ni in range_constexpr(MI_N):
-                    col = n_offset + fx.Index(wave_n * WARP_N + ni * MFMA_N + c_n)
+                    col, col_loc = _cols(ni)
                     a = Vec(acc[mi * MI_N + ni])
                     if const_expr(has_bias and not use_splitk):
-                        col_i = fx.Int32(col)
-                        if const_expr(n_tail):
-                            col_i = arith.select(col < fx.Index(k), col_i, fx.Int32(0))
-                        fx.copy(bias_atom, fx.slice(bias_div, (None, col_i)), bias_reg)
-                        bias_val = fx.Float32(fx.memref_load_vec(bias_reg)[0])
+                        bias_val = bias_vals[ni]
 
                     if const_expr(_vec_store):
-                        row0 = fx.Index(row_base)
+                        row0 = fx.Int64(row_base)
                         off_nk0 = col * dhw + row0
 
                         def _emit_vec():
@@ -636,20 +878,25 @@ def _emit_vec():
                                 vals.append(cval.to(elem_ty))
                             v4 = fx.Vector.from_elements(vals, dtype=elem_ty)
                             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))))
+                            fx.copy(y_atom_4, y_reg_4, fx.slice(y_div, (None, _route(off_nk0, row0, col_loc))))
 
-                        if const_expr(_need_chk):
-                            if _valid_raw(row0, col):
+                        if const_expr(_need_chk and not _route_store):
+                            if _valid_raw(row0, col_loc):
                                 _emit_vec()
                         else:
                             _emit_vec()
                         continue
 
                     for i in range_constexpr(MFMA_C_VALUES):
-                        row = fx.Index(row_base + i)
+                        row = fx.Int64(row_base + i)
                         off_sk = row * k + col
 
-                        if const_expr(n == 1):
+                        if const_expr(out_ndhwc):
+                            # (n, do, ho, wo, k) is row-major over exactly the GEMM's own
+                            # (npq, k) index space, so the scatter collapses to off_sk and
+                            # the row decomposition disappears.
+                            off_nk = off_sk
+                        elif const_expr(n == 1):
                             off_nk = col * dhw + row
                         else:
                             n_idx = row // dhw
@@ -667,10 +914,10 @@ def _emit():
                                     _big_store(fx.Int64(off_nk), cval)
                                 else:
                                     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))))
+                                    fx.copy(y_atom_1, y_reg_1, fx.slice(y_div, (None, _route(off_nk, row, col_loc))))
 
-                        if const_expr(_need_chk):
-                            if _valid_raw(row, col):
+                        if const_expr(_need_chk and not _route_store):
+                            if _valid_raw(row, col_loc):
                                 _emit()
                         else:
                             _emit()
@@ -680,31 +927,61 @@ def _emit():
     @flyc.jit
     def launch(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
         conv3d_implicit_kernel(y, x, weight, bias).launch(
-            grid=(grid_m, grid_n, splitk), block=(BLOCK_THREADS, 1, 1), stream=stream
+            grid=(grid_x, grid_n, m_chunks * splitk), block=(BLOCK_THREADS, 1, 1), stream=stream
         )
 
-    return launch
+    def _launch(y, x, weight, bias, stream=None):
+        with CompilationContext.compile_hints(CONV_COMPILE_HINTS):
+            return launch(y, x, weight, bias, stream=_as_stream(stream))
+
+    def _compile(y, x, weight, bias, stream=None):
+        with CompilationContext.compile_hints(CONV_COMPILE_HINTS):
+            return flyc.compile(launch, y, x, weight, bias, _as_stream(stream))
+
+    _launch.compile = _compile
+    return _launch
+
+
+SPLITK_MAX_STAGING_BYTES = 0xFFFFFFFF
+
 
+def _num_cu(device):
+    try:
+        return torch.cuda.get_device_properties(device).multi_processor_count
+    except Exception:
+        return 256
 
-def _resolve_splitk(splitk, npq, crs, k, device, tile=DEFAULT_TILE):
+
+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]]
+    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)
+        if blocks * wave_m * wave_n >= target:
+            return (tile_m, tile_n, wave_m, wave_n)
+    return legal[-1]
+
+
+def _resolve_splitk(splitk, npq, crs, k, device, tile=DEFAULT_TILE, groups=1):
     k_tiles = (crs + TILE_K - 1) // TILE_K
+    if npq * k * 4 > SPLITK_MAX_STAGING_BYTES:
+        return 1
     if splitk is None:
         tile_m, tile_n = tile[0], tile[1]
-        base = ((npq + tile_m - 1) // tile_m) * ((k + tile_n - 1) // tile_n)
+        kg = k // groups
+        base = ((npq + tile_m - 1) // tile_m) * groups * ((kg + tile_n - 1) // tile_n)
         if (
             npq < 4096
             or k_tiles < 16
-            or k % tile_n != 0
+            or kg % tile_n != 0
             or npq % tile_m != 0
             or crs % TILE_K != 0
             or npq * k * 4 > 0x7FFFFFFF
         ):
             sk = 1
         else:
-            try:
-                num_cu = torch.cuda.get_device_properties(device).multi_processor_count
-            except Exception:
-                num_cu = 256
+            num_cu = _num_cu(device)
             if base >= (3 * num_cu) // 4:
                 sk = 1
             else:
@@ -716,56 +993,193 @@ 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):
-    n, c, d, h, w = x.shape
+def _as_tuple(v, rank, name):
+    if isinstance(v, int):
+        return (v,) * rank
+    t = tuple(v)
+    if len(t) == 1:
+        return t * rank
+    assert len(t) == rank, f"{name} must be an int or a sequence of 1 or {rank} ints, got {tuple(v)}"
+    return t
+
+
+def _resolve_padding(padding, kernel, stride, dilation):
+    """Normalize torch's ``padding`` argument to a (low, high) pair of per-axis triples."""
+    if not isinstance(padding, str):
+        p = _as_tuple(padding, 3, "padding")
+        assert min(p) >= 0, f"negative padding is not supported, got (pt, ph, pw) = {p}"
+        return p, p
+    if padding == "valid":
+        return (0, 0, 0), (0, 0, 0)
+    if padding != "same":
+        raise ValueError(f"padding string must be 'same' or 'valid', got {padding!r}")
+    assert all(
+        s == 1 for s in stride
+    ), f"padding='same' is not supported for strided convolutions, got stride {tuple(stride)}"
+    total = [dl * (kn - 1) for kn, dl in zip(kernel, dilation)]
+    return tuple(t // 2 for t in total), tuple(t - t // 2 for t in total)
+
+
+def _conv3d_impl(
+    x,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    padding_mode="zeros",
+    splitk=None,
+    stream=None,
+    tile=None,
+    autotune=None,
+    input_layout="NCDHW",
+    output_layout="NCDHW",
+):
+    _check_layouts(3, input_layout, output_layout)
+
+    in_ndhwc = input_layout == "NDHWC"
+    out_ndhwc = output_layout == "NDHWC"
+    n, c, d, h, w = _shape_ncdhw(x, in_ndhwc)
     k, wc, kt, kh, kw = weight.shape
-    assert c == wc
-    assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16
-    st, sh, sw = (stride, stride, stride) if isinstance(stride, int) else stride
-    pt, ph, pw = (padding, padding, padding) if isinstance(padding, int) else padding
 
-    # 1x1x1 fast path: y[n,k,dhw] = sum_c weight[k,c] * x[n,c,dhw] — pure channel GEMM.
-    if kt == 1 and kh == 1 and kw == 1 and st == 1 and sh == 1 and sw == 1 and pt == 0 and ph == 0 and pw == 0:
+    for name, t in (("x", x), ("weight", weight), ("bias", bias)):
+        assert t is None or t.is_cuda, f"conv3d_implicit needs GPU tensors; {name} is on {t.device}"
+    assert (
+        x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16
+    ), f"conv3d_implicit is a bf16-only kernel; got x={x.dtype}, weight={weight.dtype}"
+    assert bias is None or (bias.dim() == 1 and bias.numel() == k), (
+        f"bias must be a 1-D tensor of {k} elements, one per output channel; " f"got shape {tuple(bias.shape)}"
+    )
+    groups = int(groups)
+    assert groups >= 1, f"groups must be >= 1, got {groups}"
+    assert c % groups == 0, f"in-channels {c} not divisible by groups {groups}"
+    assert k % groups == 0, f"out-channels {k} not divisible by groups {groups}"
+    assert wc == c // groups, f"weight in-channels {wc} != C/groups = {c // groups}"
+    st, sh, sw = _as_tuple(stride, 3, "stride")
+
+    assert min(st, sh, sw) >= 1, f"non-positive stride is not supported, got (st, sh, sw) = {(st, sh, sw)}"
+    dt, dh, dw = _as_tuple(dilation, 3, "dilation")
+    assert min(dt, dh, dw) >= 1, f"dilation must be >= 1, got {(dt, dh, dw)}"
+    pad_lo, pad_hi = _resolve_padding(padding, (kt, kh, kw), (st, sh, sw), (dt, dh, dw))
+    pt, ph, pw = pad_lo
+    assert padding_mode in PADDING_MODES, f"padding_mode must be one of {PADDING_MODES}, got {padding_mode!r}"
+
+    if padding_mode in ("reflect", "circular"):
+        for ax, (p, ext) in enumerate(zip(map(max, pad_lo, pad_hi), (d, h, w))):
+            if padding_mode == "reflect":
+                assert p < ext, f"reflect padding {p} must be < input extent {ext} on spatial axis {ax}"
+            else:
+                assert p <= ext, f"circular padding {p} must be <= input extent {ext} on spatial axis {ax}"
+
+    if pad_lo != pad_hi:
+        if padding_mode == "zeros":
+            x, in_ndhwc = _pad_spatial(x, in_ndhwc, (0, pad_hi[2] - pw, 0, pad_hi[1] - ph, 0, pad_hi[0] - pt))
+        else:
+            x, in_ndhwc = _pad_spatial(x, in_ndhwc, (pw, pad_hi[2], ph, pad_hi[1], pt, pad_hi[0]), mode=padding_mode)
+            pt = ph = pw = 0
+        n, c, d, h, w = _shape_ncdhw(x, in_ndhwc)
+
+    inline_pad = padding_mode != "zeros" and bool(pt or ph or pw)
+    if inline_pad and _big_in(n, c, groups, d, h, w, pt, ph, pw):
+        x, in_ndhwc = _pad_spatial(x, in_ndhwc, (pw, pw, ph, ph, pt, pt), mode=padding_mode)
+        n, c, d, h, w = _shape_ncdhw(x, in_ndhwc)
+        pt = ph = pw = 0
+        inline_pad = False
+    pad_mode = padding_mode if inline_pad else "zeros"
+
+    if (
+        groups == 1
+        and kt == 1
+        and kh == 1
+        and kw == 1
+        and st == 1
+        and sh == 1
+        and sw == 1
+        and pt == 0
+        and ph == 0
+        and pw == 0
+    ):
         wm = weight.reshape(k, c)
+        if in_ndhwc:
+            y = torch.matmul(x.reshape(n * d * h * w, c), wm.t()).reshape(n, d, h, w, k)
+            if bias is not None:
+                y = y + bias.to(y.dtype)
+            return y if out_ndhwc else y.permute(0, 4, 1, 2, 3).contiguous()
         if n == 1:
             y = torch.matmul(wm, x.reshape(c, d * h * w)).reshape(n, k, d, h, w)
         else:
             y = torch.matmul(wm, x.reshape(n, c, d * h * w)).reshape(n, k, d, h, w)
         if bias is not None:
             y = y + bias.to(y.dtype).view(1, k, 1, 1, 1)
-        return y
+        return y.permute(0, 2, 3, 4, 1).contiguous() if out_ndhwc else y
 
-    do = (d + 2 * pt - kt) // st + 1
-    ho = (h + 2 * ph - kh) // sh + 1
-    wo = (w + 2 * pw - kw) // sw + 1
+    do = (d + 2 * pt - (dt * (kt - 1) + 1)) // st + 1
+    ho = (h + 2 * ph - (dh * (kh - 1) + 1)) // sh + 1
+    wo = (w + 2 * pw - (dw * (kw - 1) + 1)) // sw + 1
+    assert min(do, ho, wo) >= 1, f"dilated filter is larger than the padded input: output ({do}, {ho}, {wo})"
     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
+    if n == 0:
+        empty = (0, do, ho, wo, k) if out_ndhwc else (0, k, do, ho, wo)
+        return torch.empty(empty, device=x.device, dtype=torch.bfloat16)
+
+    cg = c // groups
+    cgp = _pad_channels(cg)
+    if cgp != cg:
+        if in_ndhwc:
+            x = torch.nn.functional.pad(x.reshape(n, d, h, w, groups, cg), (0, cgp - cg))
+            x = x.reshape(n, d, h, w, groups * cgp)
+        else:
+            x = torch.nn.functional.pad(x.reshape(n, groups, cg, d, h, w), (0, 0, 0, 0, 0, 0, 0, cgp - cg))
+            x = x.reshape(n, groups * cgp, d, h, w)
+    c = groups * cgp
+    crs = cgp * 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)
+    x_ndhwc = x.contiguous() if in_ndhwc else _ncdhw_to_ndhwc(x, stream)
     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)
+    shape = (n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, dt, dh, dw, pad_mode, has_bias, groups, out_ndhwc)
 
     def _run(the_tile, the_wgm=1):
-        sk = _resolve_splitk(splitk, npq, crs, k, x.device, the_tile)
+        sk = _resolve_splitk(splitk, npq, crs, k, x.device, the_tile, groups)
         if sk > 1:
             y = torch.zeros((npq, k), device=x.device, dtype=torch.float32)
         else:
-            y = torch.empty((n, k, do, ho, wo), device=x.device, dtype=torch.bfloat16)
+            out_shape = (n, do, ho, wo, k) if out_ndhwc else (n, k, do, ho, wo)
+            y = torch.empty(out_shape, device=x.device, dtype=torch.bfloat16)
         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
+            n,
+            c,
+            d,
+            h,
+            w,
+            k,
+            kt,
+            kh,
+            kw,
+            st,
+            sh,
+            sw,
+            pt,
+            ph,
+            pw,
+            dt,
+            dh,
+            dw,
+            pad_mode,
+            has_bias,
+            sk,
+            the_tile,
+            the_wgm,
+            groups,
+            out_ndhwc,
         )
-        exe(y, x_ndhwc, w_packed, bias_arg, launch_stream)
+        _dispatch(exe, y, x_ndhwc, w_packed, bias_arg, stream=launch_stream)
         return y, sk
 
     if tile is not None:
@@ -778,57 +1192,143 @@ def _run(the_tile, the_wgm=1):
         best = autotune_conv3d("bf16", shape, "bf16", candidates, x.device, lambda tw: _run(tw[0], tw[1])[0])
         chosen_tile, chosen_wgm = best
     else:
-        chosen_tile = DEFAULT_TILE
+        chosen_tile = _pick_tile(npq, k, groups, x.device)
         chosen_wgm = 1
 
     y, sk = _run(chosen_tile, chosen_wgm)
     if sk > 1:
         if has_bias:
             y = y + bias_arg.view(1, k)
-        y = y.to(torch.bfloat16)
-        return y.view(n, do, ho, wo, k).permute(0, 4, 1, 2, 3)
+        if out_ndhwc:
+            return y.view(n, do, ho, wo, k).to(torch.bfloat16)
+        out = torch.empty((n, k, do, ho, wo), device=x.device, dtype=torch.bfloat16)
+        out.copy_(y.view(n, do, ho, wo, k).permute(0, 4, 1, 2, 3))
+        return out
     return y
 
 
-def _conv2d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs):
+def _conv2d_impl(
+    x, weight, bias=None, stride=1, padding=0, dilation=1, input_layout="NCHW", output_layout="NCHW", **kwargs
+):
     assert x.dim() == 4 and weight.dim() == 4, "conv2d expects (N,C,H,W) / (K,C,R,S)"
-    sh, sw = (stride, stride) if isinstance(stride, int) else stride
-    ph, pw = (padding, padding) if isinstance(padding, int) else padding
-    n, c, h, w = x.shape
+    _check_layouts(2, input_layout, output_layout)
+    sh, sw = _as_tuple(stride, 2, "stride")
+    dh, dw = _as_tuple(dilation, 2, "dilation")
+
+    if isinstance(padding, str):
+        p3 = padding
+    else:
+        ph, pw = _as_tuple(padding, 2, "padding")
+        p3 = (0, ph, pw)
     k, wc, r, s = weight.shape
-    x5 = x.reshape(n, c, 1, h, w)
+
+    if input_layout == "NHWC":
+        n, h, w, c = x.shape
+        x5, in5 = x.reshape(n, 1, h, w, c), "NDHWC"
+    else:
+        n, c, h, w = x.shape
+        x5, in5 = x.reshape(n, c, 1, h, w), "NCDHW"
+    out5 = "NDHWC" if output_layout == "NHWC" else "NCDHW"
     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=p3,
+        dilation=(1, dh, dw),
+        input_layout=in5,
+        output_layout=out5,
+        **kwargs,
+    )
+    if output_layout == "NHWC":
+        return y5.reshape(y5.shape[0], y5.shape[2], y5.shape[3], y5.shape[4])
     return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[3], y5.shape[4])
 
 
-def _conv1d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs):
+def _conv1d_impl(
+    x, weight, bias=None, stride=1, padding=0, dilation=1, input_layout="NCW", output_layout="NCW", **kwargs
+):
     assert x.dim() == 3 and weight.dim() == 3, "conv1d expects (N,C,W) / (K,C,S)"
-    sw = stride if isinstance(stride, int) else stride[0]
-    pw = padding if isinstance(padding, int) else padding[0]
-    n, c, w = x.shape
+    _check_layouts(1, input_layout, output_layout)
+    (sw,) = _as_tuple(stride, 1, "stride")
+    (dw,) = _as_tuple(dilation, 1, "dilation")
+    if isinstance(padding, str):
+        p3 = padding
+    else:
+        p3 = (0, 0, _as_tuple(padding, 1, "padding")[0])
     k, wc, s = weight.shape
-    x5 = x.reshape(n, c, 1, 1, w)
+    if input_layout == "NWC":
+        n, w, c = x.shape
+        x5, in5 = x.reshape(n, 1, 1, w, c), "NDHWC"
+    else:
+        n, c, w = x.shape
+        x5, in5 = x.reshape(n, c, 1, 1, w), "NCDHW"
+    out5 = "NDHWC" if output_layout == "NWC" else "NCDHW"
     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=p3,
+        dilation=(1, 1, dw),
+        input_layout=in5,
+        output_layout=out5,
+        **kwargs,
+    )
+    if output_layout == "NWC":
+        return y5.reshape(y5.shape[0], y5.shape[3], y5.shape[4])
     return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[4])
 
 
-def conv3d_implicit(x, weight, bias=None, stride=1, padding=0, **kwargs):
+def conv3d_implicit(x, weight, bias=None, stride=1, padding=0, dilation=1, **kwargs):
     """Main implicit-GEMM conv entry; dispatches 1D/2D/3D by filter rank.
 
-    Rank is taken from the filter (weight.dim() - 2): 3 -> 3D (N,C,D,H,W)/(K,C,T,R,S),
-    2 -> 2D (N,C,H,W)/(K,C,R,S), 1 -> 1D (N,C,W)/(K,C,S); x and weight must match.
-    True 3D calls run the implementation directly; 2D/1D reshape to the degenerate
-    5D case. stride/padding/bias and extra kwargs (splitk, tile, autotune, stream)
-    forward to the chosen path.
+    Rank is taken from the filter (weight.dim() - 2): 3 -> 3D (N,C,D,H,W)/(K,C,T,R,S).
+
+    ``input_layout`` and ``output_layout`` are independent and named per rank:
+    "NCDHW"/"NDHWC", "NCHW"/"NHWC", "NCW"/"NWC". The weight stays KC*, and the batch axis
+    leads in both, so an unbatched input works either way. Channels-last is the kernel's
+    own layout on both sides: an NDHWC input skips the pre-transpose, and an NDHWC output
+    is the (npq, K) index space the GEMM already writes, so it also skips the split-K
+    epilogue's transpose. Channels-last output does give up the vectorized store on the
+    ``n == 1`` fast path, since a lane's four accumulator values are four M rows and those
+    are K apart once channels are innermost.
+
+    ``padding`` takes an int, a per-axis tuple, or one of torch's two strings. "valid" is
+    no padding. "same" pads so the output keeps the input's spatial extent, which needs
+    ``dilation * (kernel - 1)`` elements per axis and, like torch, is only defined at
+    stride 1. That total is normally even and costs nothing beyond an ordinary symmetric
+    pad. An even-length filter under odd dilation makes it odd, and torch's rule of
+    putting the extra element on the high side then asks for a pad the kernel cannot
+    express with one value per axis; that case materializes a padded input first, exactly
+    as torch does (it warns about the same copy). ``padding_mode`` applies to "same" too.
+
+    ``dilation`` follows torch semantics: it spaces the filter taps by that factor
+    over the input, shrinking the output to
+    ``(D + 2*pad - dilation*(T-1) - 1)//stride + 1`` per axis. It costs nothing in the
+    GEMM -- the K axis is still C/groups*T*R*S -- it only stretches the im2col gather,
+    so a dilated filter reads a wider input footprint per output row and gets less
+    reuse out of cache than the same filter undilated.
+
+    ``groups`` follows torch semantics: C and K must both be divisible by it and the
+    weight's channel dim is C/groups. Groups map onto the N grid axis, one tile never
+    spanning two groups, so efficiency tracks how well K/groups fills TILE_N. Measured
+    on gfx950 vs torch/MIOpen, moderate cardinality wins across the board (1.5-2.0x for
+    K/groups in [8, 256]). True depthwise (groups == C, so C/groups == 1) is the one
+    weak case at ~0.5x: C/groups=1 pads to the gather's 8-wide vector, wasting 7/8 of
+    the K axis, while K/groups=1 leaves all but one column of the N tile masked.
+    Narrower tiles recover little there -- depthwise wants its own kernel, not this
+    single-GEMM mapping.
     """
-    assert x.dim() == weight.dim(), f"x rank {x.dim()} != weight rank {weight.dim()}"
     spatial_rank = weight.dim() - 2
-    if spatial_rank == 3:
-        return _conv3d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs)
-    if spatial_rank == 2:
-        return _conv2d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs)
-    if spatial_rank == 1:
-        return _conv1d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs)
-    raise ValueError(f"conv3d_implicit supports 1D/2D/3D; got filter rank {weight.dim()}")
+    if spatial_rank not in (1, 2, 3):
+        raise ValueError(f"conv3d_implicit supports 1D/2D/3D; got filter rank {weight.dim()}")
+    unbatched = x.dim() == weight.dim() - 1
+    if unbatched:
+        x = x.unsqueeze(0)
+    assert x.dim() == weight.dim(), f"x rank {x.dim()} != weight rank {weight.dim()}"
+    impl = {3: _conv3d_impl, 2: _conv2d_impl, 1: _conv1d_impl}[spatial_rank]
+    y = impl(x, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, **kwargs)
+    return y.squeeze(0) if unbatched else y
diff --git a/kernels/conv/conv3d_implicit_fp8.py b/kernels/conv/conv3d_implicit_fp8.py
index 7ac3115091..85515ab450 100644
--- a/kernels/conv/conv3d_implicit_fp8.py
+++ b/kernels/conv/conv3d_implicit_fp8.py
@@ -172,11 +172,11 @@ def lds_load_scalar(elem_offset):
             sv = (lin % TR_VPL) * TR_VEC
             cc = c0 + rc
             ss = s0 + sv
-            valid = (cc < fx.Index(c)) & (ss < fx.Index(s))
+            valid = (cc < fx.Int64(c)) & (ss < fx.Int64(s))
             if const_expr(TR_BIG):
                 # Clamp OOB coords to 0 so the raw load never dereferences past the tensor.
-                cc_s = fx.Index(arith.select(valid, fx.Int64(cc), fx.Int64(0)))
-                ss_s = fx.Index(arith.select(valid, fx.Int64(ss), fx.Int64(0)))
+                cc_s = fx.Int64(arith.select(valid, fx.Int64(cc), fx.Int64(0)))
+                ss_s = fx.Int64(arith.select(valid, fx.Int64(ss), fx.Int64(0)))
                 addr = in_base_addr + (fx.Int64(nb) * fx.Int64(c) + fx.Int64(cc_s)) * fx.Int64(s) + fx.Int64(ss_s)
                 ptr = _global_ptr_from_addr(addr, u8.ir_type, inp).llvm_ptr
                 v = llvm.LoadOp(fx.Vector.make_type(4, fx.Int32), ptr, alignment=16).result
@@ -197,7 +197,7 @@ def lds_load_scalar(elem_offset):
             cv = (lin % TR_VPL) * TR_VEC
             ss = s0 + rs
             cc = c0 + cv
-            valid = (ss < fx.Index(s)) & (cc < fx.Index(c))
+            valid = (ss < fx.Int64(s)) & (cc < fx.Int64(c))
             if valid:
                 scalars = [lds_load_scalar((cv + j) * TR_LDS_S + rs) for j in range_constexpr(TR_VEC)]
                 packed_u8 = fx.Vector.from_elements(scalars, dtype=u8)
@@ -348,15 +348,15 @@ def conv3d_implicit_fp8_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bi
         if const_expr(WGM > 1):
             # Grouped-M L2-swizzle: visit WGM consecutive m-tiles across all n-tiles
             # before advancing, keeping the weight tile hot in L2. Mirrors the bf16 kernel.
-            pid = fx.Index(fx.block_idx.x) + fx.Index(fx.block_idx.y) * fx.Index(grid_m)
-            blocks_per_group = fx.Index(WGM * grid_n)
+            pid = fx.Int64(fx.block_idx.x) + fx.Int64(fx.block_idx.y) * fx.Int64(grid_m)
+            blocks_per_group = fx.Int64(WGM * grid_n)
             group_id = pid // blocks_per_group
-            first_m = group_id * fx.Index(WGM)
-            group_rows = fx.Index(grid_m) - first_m
-            group_rows = fx.Index(arith.select(group_rows < fx.Index(WGM), group_rows, fx.Index(WGM)))
+            first_m = group_id * fx.Int64(WGM)
+            group_rows = fx.Int64(grid_m) - first_m
+            group_rows = fx.Int64(arith.select(group_rows < fx.Int64(WGM), group_rows, fx.Int64(WGM)))
             local = pid % blocks_per_group
-            block_m = fx.Index(first_m + (local % group_rows))
-            block_n = fx.Index(local // group_rows)
+            block_m = fx.Int64(first_m + (local % group_rows))
+            block_n = fx.Int64(local // group_rows)
         else:
             block_m = fx.block_idx.x
             block_n = fx.block_idx.y
@@ -368,9 +368,9 @@ def conv3d_implicit_fp8_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bi
         if const_expr(BIG_IN):
             nbase = m_offset // dhw
             ot_base0 = (m_offset % dhw) // hw_o
-            base_t = ot_base0 - fx.Index(pt)
-            base_t = arith.select(base_t < fx.Index(0), fx.Index(0), base_t)
-            x_base_byte = ((nbase * fx.Index(d) + base_t) * fx.Index(h)) * fx.Index(width) * fx.Index(c)
+            base_t = ot_base0 - fx.Int64(pt)
+            base_t = arith.select(base_t < fx.Int64(0), fx.Int64(0), base_t)
+            x_base_byte = ((nbase * fx.Int64(d) + base_t) * fx.Int64(h)) * fx.Int64(width) * fx.Int64(c)
             x_addr = fx.Int64(fx.ptrtoint(fx.get_iter(x))) + fx.Int64(x_base_byte)
             x_div = fx.logical_divide(
                 _make_fp8_buffer_tensor_from_addr(x_addr, f8_ir_t, x_buf),
@@ -378,7 +378,7 @@ def conv3d_implicit_fp8_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bi
             )
 
         def in_range(v, hi):
-            return (v >= fx.Index(0)) & (v < fx.Index(hi))
+            return (v >= fx.Int64(0)) & (v < fx.Int64(hi))
 
         # ---- im2col address for a (M-row, K-col) chunk of 16 contiguous channels ----
         # k_col is 16-aligned and c % 16 == 0, so the 16 elements at k_col are consecutive
@@ -387,7 +387,7 @@ def in_range(v, hi):
         # The output-spatial decomposition of m_row is loop-invariant across the K loop
         # but costs 4 div/mods, so it is precomputed once per (half, step) and reused.
         def _spatial_of_row(m_row):
-            row_valid = m_row < fx.Index(npq)
+            row_valid = m_row < fx.Int64(npq)
             if const_expr(temporal_only_fast):
                 out_t = (m_row // hw_o) % d
                 return (row_valid, out_t, m_row)
@@ -407,7 +407,7 @@ def im2col_safe_elem_pre(spatial, k_col, k_iter):
             # tap), so gate on k_iter and mask those to the OOB sentinel.
             k_in_range = True
             if const_expr(K_PARTIAL and k_iter == K_ITERS - 1):
-                k_in_range = k_col < fx.Index(crs)
+                k_in_range = k_col < fx.Int64(crs)
             if const_expr(temporal_only_fast):
                 row_valid, out_t, m_row = spatial
                 kt_i = k_col // c
@@ -461,7 +461,7 @@ def im2col_safe_elem_pre(spatial, k_col, k_iter):
         # Loop-invariant per (half, step): LDS byte offset and m_row's spatial decomposition.
         # Computed once here so the SSA values dominate every conv_a_g2s call.
         def _conv_a_g2s_pre(half):
-            m_half_base = m_offset + fx.Index(half * LDS_BLOCK_M)
+            m_half_base = m_offset + fx.Int64(half * LDS_BLOCK_M)
             steps = []
             for step in range_constexpr(N_LDS_STEPS_A):
                 row_g = lane_id // 8 + wave_id * 8 + step * (N_WAVES * 8)
@@ -478,12 +478,12 @@ def _conv_a_g2s_pre(half):
         _a_g2s_pre = [_conv_a_g2s_pre(0), _conv_a_g2s_pre(1)]
 
         def conv_a_g2s(lds_dst, half, k_iter):
-            k_base = fx.Index(k_iter * BLOCK_K)
+            k_base = fx.Int64(k_iter * BLOCK_K)
             for step_off, cc, m_row, spatial in _a_g2s_pre[half]:
                 if const_expr(STRIP_IM2COL):
                     k_col = k_base + cc
                     lin = m_row * crs + k_col
-                    safe = fx.Int32(arith.select(m_row < fx.Index(npq), lin, fx.Index(OOB_SENTINEL_ELEM)))
+                    safe = fx.Int32(arith.select(m_row < fx.Int64(npq), lin, fx.Int64(OOB_SENTINEL_ELEM)))
                 else:
                     safe = im2col_safe_elem_pre(spatial, k_base + cc, k_iter)
                 base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + fx.Int32(step_off)
@@ -623,18 +623,18 @@ def _big_store(off_elem, value):
         def store_cfrag(c_frag, base_row, base_col):
             for ti in range_constexpr(N_TILES_A):
                 for tj in range_constexpr(N_TILES_B):
-                    col = base_col + fx.Index(tj * 16) + lane_id % 16
-                    col_valid = col < fx.Index(k)
+                    col = base_col + fx.Int64(tj * 16) + lane_id % 16
+                    col_valid = col < fx.Int64(k)
                     if const_expr(has_bias):
-                        col_i = fx.Int32(arith.select(col_valid, col, fx.Index(0)))
+                        col_i = fx.Int32(arith.select(col_valid, col, fx.Int64(0)))
                         fx.copy(bias_atom, fx.slice(bias_div, (None, col_i)), bias_reg)
                         bias_val = fx.Float32(fx.memref_load_vec(bias_reg)[0])
                     vec_f32 = Vec(c_frag[mfma.idx(ti, tj)])
-                    row_base = base_row + fx.Index(ti * 16) + (lane_id // 16) * 4
+                    row_base = base_row + fx.Int64(ti * 16) + (lane_id // 16) * 4
 
                     if const_expr(_vec_store):
                         off0 = col * dhw + row_base
-                        row_ok = col_valid & (row_base + fx.Index(3) < fx.Index(npq))
+                        row_ok = col_valid & (row_base + fx.Int64(3) < fx.Int64(npq))
                         if row_ok:
                             vals = []
                             for i in range_constexpr(4):
@@ -646,11 +646,11 @@ def store_cfrag(c_frag, base_row, base_col):
                         continue
 
                     for i in range_constexpr(4):
-                        row = row_base + fx.Index(i)
+                        row = row_base + fx.Int64(i)
                         out = vec_f32[i]
                         if const_expr(has_bias):
                             out = out + bias_val
-                        valid = col_valid & (row < fx.Index(npq))
+                        valid = col_valid & (row < fx.Int64(npq))
                         if const_expr(n == 1):
                             off_ncdhw = col * dhw + row
                         else:
@@ -663,19 +663,19 @@ def store_cfrag(c_frag, base_row, base_col):
                         else:
                             # Route masked-off lanes one element past the end; the
                             # descriptor's num_records bound drops them in hardware.
-                            off_i = fx.Int32(arith.select(valid, off_ncdhw, fx.Index(npq * k)))
+                            off_i = fx.Int32(arith.select(valid, off_ncdhw, fx.Int64(npq * k)))
                             fx.memref_store_vec(Vec.filled(1, out.to(fx.BFloat16), fx.BFloat16), y_reg_1)
                             fx.copy(y_atom_1, y_reg_1, fx.slice(y_div, (None, off_i)))
 
         wave_m_offset = wave_m * (N_TILES_A * 16)
         wave_n_offset = wave_n * (N_TILES_B * 16)
-        base_row = m_offset + fx.Index(wave_m_offset)
-        base_col = block_n * BLOCK_N + fx.Index(wave_n_offset)
+        base_row = m_offset + fx.Int64(wave_m_offset)
+        base_col = block_n * BLOCK_N + fx.Int64(wave_n_offset)
 
         store_cfrag(c00_frag, base_row, base_col)
-        store_cfrag(c01_frag, base_row, base_col + fx.Index(LDS_BLOCK_N))
-        store_cfrag(c10_frag, base_row + fx.Index(LDS_BLOCK_M), base_col)
-        store_cfrag(c11_frag, base_row + fx.Index(LDS_BLOCK_M), base_col + fx.Index(LDS_BLOCK_N))
+        store_cfrag(c01_frag, base_row, base_col + fx.Int64(LDS_BLOCK_N))
+        store_cfrag(c10_frag, base_row + fx.Int64(LDS_BLOCK_M), base_col)
+        store_cfrag(c11_frag, base_row + fx.Int64(LDS_BLOCK_M), base_col + fx.Int64(LDS_BLOCK_N))
 
     @flyc.jit
     def launch(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
diff --git a/kernels/gemm/fp4_gemm_4wave.py b/kernels/gemm/fp4_gemm_4wave.py
index 70b2a23342..93d776fd75 100644
--- a/kernels/gemm/fp4_gemm_4wave.py
+++ b/kernels/gemm/fp4_gemm_4wave.py
@@ -495,10 +495,14 @@ def __init__(self, a_scale, b_scale, K, lane_id, wave_id, lds_base_ptr, a_rows,
         # elements, i.e. K//32 bytes per row.
         row_bytes = K // 32
         self.a_rsrc = fx.as_ir_value(
-            _buffer_ops.create_buffer_resource(a_scale, max_size=False, num_records_bytes=a_rows * row_bytes)
+            fx.rocdl.get_buffer_rsrc(
+                _buffer_ops.create_buffer_resource(a_scale, max_size=False, num_records_bytes=a_rows * row_bytes)
+            )
         )
         self.b_rsrc = fx.as_ir_value(
-            _buffer_ops.create_buffer_resource(b_scale, max_size=False, num_records_bytes=b_rows * row_bytes)
+            fx.rocdl.get_buffer_rsrc(
+                _buffer_ops.create_buffer_resource(b_scale, max_size=False, num_records_bytes=b_rows * row_bytes)
+            )
         )
         # Per-lane block / within-block index (loop-invariant).
         self._blk = lane_id // 16  # 0..3 -> which of the 4 blocks
@@ -751,8 +755,12 @@ def _gather_scale_thunks(k, slot):
         gl_off_a = _global_swizzle(lane_id, wave_id, K_BYTES, N_LDS_ROUNDS, False)
         gl_off_b = _global_swizzle(lane_id, wave_id, K_BYTES, N_LDS_ROUNDS, True)
 
-        a_rsrc = _buffer_ops.create_buffer_resource(A, max_size=False, num_records_bytes=c_m * K_BYTES)
-        b_rsrc = _buffer_ops.create_buffer_resource(B_T, max_size=False, num_records_bytes=c_n * K_BYTES)
+        a_rsrc = fx.rocdl.get_buffer_rsrc(
+            _buffer_ops.create_buffer_resource(A, max_size=False, num_records_bytes=c_m * K_BYTES)
+        )
+        b_rsrc = fx.rocdl.get_buffer_rsrc(
+            _buffer_ops.create_buffer_resource(B_T, max_size=False, num_records_bytes=c_n * K_BYTES)
+        )
         a_g2s = G2SLoaderAsm(a_rsrc, gl_off_a, N_TILES_A, wave_id, scopes=_sc, base_ptr=_base_ptr)
         b_g2s = G2SLoaderAsm(b_rsrc, gl_off_b, N_TILES_B, wave_id, scopes=_sc, base_ptr=_base_ptr)
         # Precompute the g2s wave-uniform LDS base into SGPR once (all 8 buffers share
diff --git a/kernels/moe/moe_2stage_a16wmix/gemm1.py b/kernels/moe/moe_2stage_a16wmix/gemm1.py
index d04beac669..3cdbb916b3 100644
--- a/kernels/moe/moe_2stage_a16wmix/gemm1.py
+++ b/kernels/moe/moe_2stage_a16wmix/gemm1.py
@@ -696,7 +696,7 @@ def _reduce_round(accs):
                     y = _silu_mul_batch([g], [u])[0]
                 yb = y.to(fx.BFloat16)
                 out_idx = sorted_row * inter_i32 + col_g_list[ni]
-                buffer_ops.buffer_store(yb, _raw(out_rsrc), _raw(out_idx), mask=valid)
+                buffer_ops.buffer_store(yb, out_rsrc, _raw(out_idx), mask=valid)
 
 
 def gemm1_a16w4_grid(BM, *, INTER, TILE_N, max_m_blocks):