diff --git a/ethosu/vela/api.py b/ethosu/vela/api.py
index e91c0bd..e31c373 100644
--- a/ethosu/vela/api.py
+++ b/ethosu/vela/api.py
@@ -447,9 +447,87 @@
     This function can be used to find a valid value for npu_op.block_config.
     The block config is the unit of work in which the NPU generates the OFM.
     """
-    from . import register_command_stream_generator
+    from .architecture_features import Accelerator
+    from .architecture_features import ArchitectureFeatures
+    from .architecture_features import Block
+    from .architecture_features import create_default_arch
+    from .architecture_allocator import try_block_config
+    from .register_command_stream_generator import resampling_mode_map
+    from .register_command_stream_util import to_kernel
+    from .operation import NpuBlockType
 
-    return register_command_stream_generator.find_block_configs(npu_op, accelerator)
+    is_partkernel = False
+    if isinstance(npu_op, NpuConv2DOperation):
+        block_type = NpuBlockType.ConvolutionMxN
+        is_partkernel = npu_op.block_traversal == NpuBlockTraversal.PART_KERNEL_FIRST
+    elif isinstance(npu_op, NpuConvDepthWiseOperation):
+        block_type = NpuBlockType.ConvolutionDepthWise
+    elif isinstance(npu_op, NpuPoolingOperation):
+        block_type = NpuBlockType.ReduceSum if npu_op.sub_op_type == NpuPoolingOp.REDUCE_SUM else NpuBlockType.Pooling
+    elif isinstance(npu_op, NpuElementWiseOperation):
+        block_type = NpuBlockType.ElementWise
+    else:
+        assert 0, "Unsupported operation"
+
+    ifm_shape = Block(npu_op.ifm.shape.width, npu_op.ifm.shape.height, npu_op.ifm.shape.depth)
+    ifm2_shape = None
+    if npu_op.ifm2:
+        ifm2_shape = Block(npu_op.ifm2.shape.width, npu_op.ifm2.shape.height, npu_op.ifm2.shape.depth)
+    ofm_shape = Block(npu_op.ofm.shape.width, npu_op.ofm.shape.height, npu_op.ofm.shape.depth)
+
+    ifm_resampling_mode = resampling_mode_map[npu_op.ifm_upscale]
+    ifm_bits = npu_op.ifm.data_type.size_in_bits()
+    kernel = to_kernel(npu_op.kernel)
+    lut_banks = 0
+    if npu_op.activation:
+        lut_banks = 2 if npu_op.activation.op_type == NpuActivationOp.TABLE_LOOKUP else 0
+
+    has_scaling = True
+    for tensor in [npu_op.ifm, npu_op.ifm2, npu_op.ofm]:
+        if tensor and tensor.quantization is None:
+            has_scaling = False
+            break
+
+    arch = create_default_arch(Accelerator.from_npu_accelerator(accelerator))
+
+    max_block_width = min(arch.ofm_block_max.width, ofm_shape.width)
+    max_block_height = min(arch.ofm_block_max.height, ofm_shape.height)
+    max_block_depth = min(arch.ofm_block_max.depth, ofm_shape.depth)
+
+    min_block_height = max(arch.ofm_ublock.height, 2 if ifm_resampling_mode != NpuResamplingMode.NONE else 1)
+    min_block_width = max(arch.ofm_ublock.width, 2 if ifm_resampling_mode != NpuResamplingMode.NONE else 1)
+
+    valid_block_configs = []
+    for w in range(min_block_width, max_block_width + min_block_width, min_block_width):
+        for h in range(min_block_height, max_block_height + min_block_height, min_block_height):
+            # Try valid OFM block depths
+            for c in range(arch.ofm_ublock.depth, max_block_depth + arch.ofm_ublock.depth, arch.ofm_ublock.depth):
+                # OFM block depth has the constraint that if it causes the OFM to be
+                # split, it must be a multiple of the OFM split size
+                if (c >= max_block_depth) or (c < max_block_depth and (c % ArchitectureFeatures.OFMSplitDepth) == 0):
+                    block = Block(w, h, c)
+                    config = try_block_config(
+                        block,
+                        arch,
+                        block_type,
+                        ofm_shape,
+                        ifm_shape,
+                        ifm2_shape,
+                        npu_op.ifm2_scalar is not None,
+                        ifm_bits,
+                        is_partkernel,
+                        kernel,
+                        lut_banks,
+                        has_scaling,
+                        ifm_resampling_mode,
+                    )
+
+                    if config:
+                        ofm_block = config.ofm_block
+                        valid_block_configs.append(NpuShape3D(ofm_block.height, ofm_block.width, ofm_block.depth))
+
+    assert len(valid_block_configs) > 0
+    return valid_block_configs
 
 
 def npu_generate_register_command_stream(npu_op_list: List[NpuOperation], accelerator: NpuAccelerator) -> List[int]:
diff --git a/ethosu/vela/architecture_allocator.py b/ethosu/vela/architecture_allocator.py
index 3c49eb1..32502e3 100644
--- a/ethosu/vela/architecture_allocator.py
+++ b/ethosu/vela/architecture_allocator.py
@@ -412,4 +412,5 @@
     config.layout = layout
     config.bank_size = arch.shram_bank_size
     config.ifm_block = ifm_block
+    config.ofm_block = block_config
     return config
diff --git a/ethosu/vela/register_command_stream_generator.py b/ethosu/vela/register_command_stream_generator.py
index b4a633e..d61e571 100644
--- a/ethosu/vela/register_command_stream_generator.py
+++ b/ethosu/vela/register_command_stream_generator.py
@@ -1025,23 +1025,6 @@
     return res
 
 
-# -------------------------------------------------------------------
-# EXTERNAL API
-# -------------------------------------------------------------------
-
-
-def find_block_configs(npu_op: NpuOperation, npu_accelerator: NpuAccelerator) -> List[NpuShape3D]:
-    """
-    Internal implementation of the public facing API for finding block configs.
-    """
-    if isinstance(npu_op, NpuBlockOperation):
-        # TODO: implement this function
-        arch = create_default_arch(Accelerator.from_npu_accelerator(npu_accelerator))
-        block = arch.ofm_ublock
-        return [NpuShape3D(height=block.height, width=block.width, depth=block.depth)]
-    return []
-
-
 def generate_register_command_stream(npu_op_list: List[NpuOperation], npu_accelerator: NpuAccelerator) -> List[int]:
     """
     Internal implementation of the public facing API for generating an Ethos-U register command stream.