def _check_layernorm_n_match(op1: Operator, op2: Operator) -> int:
    """n is the flattened normalized_shape dim"""
    n1 = _get_layernorm_flattened_normalized_shape(op1)
    n2 = _get_layernorm_flattened_normalized_shape(op2)
    # we can't use half4 kernel here anyways
    if n1 < 128 and n2 < 128:
        return True

    align1 = _get_layernorm_alignment(n1)
    align2 = _get_layernorm_alignment(n2)
    if n1 <= 4096 and n2 <= 4096:
        if align1 != align2:
            return False

    # TODO: may need better heuristics
    # We may group 128, 256, 512, 1024 together if we only check neighbors, which might be ok.
    # We mostly want to rule out 128, 128, ..., 1024 cases
    if n1 >= 4 * n2 or 4 * n1 <= n2:
        return False

    if align1 == 4:

        def _in_range(n):
            return n >= 128 and n <= 4096

        # prefer to use half4 kernel
        n1_in_range = _in_range(n1)
        n2_in_range = _in_range(n2)
        return not (n1_in_range ^ n2_in_range)
    return True
n is the flattened normalized_shape dim
python/aitemplate/compiler/transform/fuse_group_ops.py