def _fuse_gemm_ops(
    op_group: List[Operator], sorted_graph: List[Tensor]
) -> List[Tensor]:
    """
    Replace a group of ops with a single group op
    """

    assert op_group[0]._attrs["op"].startswith("gemm_rcr"), (
        f"_fuse_gemm_ops only supports gemm_rcr family ops. "
        f"{op_group[0]._attrs['op']} is not supported"
    )

    # Sort by weight name
    op_group.sort(key=lambda x: x._attrs["inputs"][1]._attrs["name"])

    # Make the order deterministic, important for cache hit.
    # sort ops by decreasing K, N
    op_group.sort(
        key=lambda x: (
            x._attrs["inputs"][1]._size(1).value(),
            x._attrs["inputs"][1]._size(0).value(),
        ),
        reverse=True,
    )

    # gather inputs
    group_inputs = [op._attrs["inputs"] for op in op_group]

    # remove dst_ops from inputs
    for op in op_group:
        for input in op._attrs["inputs"]:
            transform_utils.remove_dst_op_from_tensor(input, op)

    # create group op
    op_type = op_group[0]._attrs["op"]
    assert op_type in _group_gemm_op_mapping, f"{op_type} not in _group_gemm_op_mapping"
    group_op = _group_gemm_op_mapping[op_type]

    group_outputs = group_op()(group_inputs)

    for i, op in enumerate(op_group):
        new_output = group_outputs[i]
        op_output = op._attrs["outputs"][0]
        transform_utils.replace_tensor(op_output, new_output)

    # sorted_graph is no longer sorted here
    sorted_graph.extend(group_outputs)
    return sorted_graph

    Replace a group of ops with a single group op
    
python/aitemplate/compiler/transform/fuse_group_ops.py