def _group_split_outputs_together(
    sorted_graph: List[Tensor], sorted_ops: List[Operator], op_type: str
) -> List[List[Operator]]:
    """As long as alignment allows, we group all output gemm ops from split op
    together to eliminate the cost of split. Here we don't exclude large gemms
    because the copy of split is worse than group gemm overhead.
    """
    groups = []
    if op_type not in _group_gemm_op_mapping:
        return groups
    for op in sorted_ops:
        if op._attrs["op"] != "split":
            continue
        if not _can_fuse_split_op(op):
            continue

        gemm_group = []
        for output in op._attrs["outputs"]:
            dst_ops = list(output.dst_ops())
            if len(dst_ops) == 0:
                break
            gemm_op = dst_ops[0]
            if gemm_op._attrs["op"] != op_type:
                break
            if not _filter_gemm_op(gemm_op):
                break
            if not gemm_group:
                gemm_group.append(gemm_op)
            else:
                if _gemm_op_check(gemm_group[-1], gemm_op):
                    gemm_group.append(gemm_op)
                else:
                    break
        if len(gemm_group) == len(op._attrs["outputs"]) and all(
            not _has_cycle(grouped_op, gemm_group) for grouped_op in gemm_group
        ):
            _fuse_gemm_ops(gemm_group, sorted_graph)
            groups.append(gemm_group)
    return groups
As long as alignment allows, we group all output gemm ops from split op
together to eliminate the cost of split. Here we don't exclude large gemms
because the copy of split is worse than group gemm overhead.
python/aitemplate/compiler/transform/fuse_group_ops.py