def split_large_split_ops(sorted_graph: List[Tensor], _: str) -> List[Tensor]:
    """
    Our split CUDA kernel takes an output meta argument whose size
    is proportional to the number of outputs. In extreme cases, the total size
    of the params of a split kernel may exceed the limit imposed by the CUDA
    compiler. In such cases, we split the split op into separate ones.
    """
    modified = False
    sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
    for op in sorted_ops:
        if not op._attrs["op"].startswith("split"):
            continue
        split_op = op

        split_params_size = _split_kernel_single_input_output_param_size(split_op)
        if split_params_size > MAX_CUDA_PARAM_BYTES:
            raise RuntimeError(
                f"cannot handle cases: {split_params_size=} > {MAX_CUDA_PARAM_BYTES=}"
            )
        if split_params_size * len(split_op._attrs["outputs"]) <= MAX_CUDA_PARAM_BYTES:
            continue

        modified = True
        split_dim = split_op._attrs["split_dim"]
        split_sizes = split_op._attrs["split_sizes"]
        outputs = split_op._attrs["outputs"]
        num_outputs_per_split = MAX_CUDA_PARAM_BYTES // split_params_size
        # compute how many split ops we need to fix within MAX_CUDA_PARAM_BYTES
        num_split_ops = (
            len(outputs) + num_outputs_per_split - 1
        ) // num_outputs_per_split

        output_mapping = []
        for split_i in range(num_split_ops):
            start = split_i * num_outputs_per_split
            end = min(
                (split_i + 1) * num_outputs_per_split, len(split_op._attrs["outputs"])
            )

            remove_indices = list(range(start)) + list(
                range(end, len(split_op._attrs["outputs"]))
            )
            new_split = ops.split()
            new_outputs = new_split(
                split_op._attrs["inputs"][0], split_sizes, split_dim
            )
            new_split.remove_output_at(remove_indices)
            new_outputs = new_split._attrs["outputs"]
            sorted_graph += list(new_outputs)
            output_mapping += list(zip(outputs[start:end], new_outputs))

        for old_output, new_output in output_mapping:
            transform_utils.replace_tensor(old_output, new_output)

    if not modified:
        return sorted_graph

    new_output_tensors = [
        tensor for tensor in sorted_graph if tensor._attrs["is_output"]
    ]
    sorted_graph = toposort.toposort(new_output_tensors)
    return sorted_graph
Our split CUDA kernel takes an output meta argument whose size
is proportional to the number of outputs. In extreme cases, the total size
of the params of a split kernel may exceed the limit imposed by the CUDA
compiler. In such cases, we split the split op into separate ones.
python/aitemplate/compiler/transform/split_large_split_ops.py