def process_args():
    """
    Parses command-line arguments and returns function configs and input configs.
    """
    # create parser
    parser = argparse.ArgumentParser(
        prog="Benchmark FlashAttention",
        allow_abbrev=False,
    )
    # functions
    parser.add_argument(
        "-benchmark_fn",
        type=str,
        nargs="*",
        choices=FUNCTIONS,
        required=True,
        help=f"Function(s) to benchmark",
    )
    parser.add_argument(
        "--mode",
        type=str,
        nargs='*',
        choices=VALID_MODES,
        default=None,
        help=f"Benchmarking mode(s) to run. If omitted, runs all supported modes for each function.",
    )
    # config
    parser.add_argument("-b", type=int, default=None, help="Batch size")
    parser.add_argument("-hq", type=int, default=None, help="Q Number of heads")
    parser.add_argument("-hk", type=int, default=None, help="K and V Number of heads")
    parser.add_argument("-sq", type=int, default=None, help="Q Sequence Length")
    parser.add_argument("-sk", type=int, default=None, help="K and V Sequence Length")
    parser.add_argument("-d", type=int, default=None, help="Head Dimension")
    parser.add_argument("-causal", action="store_true", default=None, help="Causal")
    parser.add_argument("-dropout", type=float, default=None, help="Dropout")

    # parse args
    args = parser.parse_args()

    # parse function args
    benchmark_fns = args.benchmark_fn
    requested_modes = args.mode 

    # fenerate function configurations and input configurations separately
    all_function_configs = []
    all_input_configs = {}  # Maps function config -> input configs
    for fn_name in benchmark_fns:
        is_varlen, is_fp8, packing, supported_dtypes, supported_backends, supported_modes_for_fn, supported_env_configs, device = get_fn_params(fn_name)
        
        # Generate or use custom input configurations
        if args.b or args.hq or args.hk or args.sq or args.sk or args.d:
            assert args.b and args.hq and args.sq and args.d, (
                "if custom config is specified, please provide at least batch, number of Q heads, Q sequence length, and head size."
            )
            
            batch = args.b
            hq = args.hq
            hk = args.hk if args.hk is not None else args.hq
            sq = args.sq
            sk = args.sk if args.sk is not None else args.sq
            d_head = args.d
            causal = args.causal if args.causal is not None else False
            dropout = args.dropout if args.dropout is not None else 0.0
            input_configs = [(batch, hq, hk, sq, sk, d_head, causal, dropout)]
        else:
            if True:
                input_configs = get_input_config_set("llama")
            else:
                input_configs = generate_benchmark_configs(is_varlen, packing)

        # filter by mode
        modes_to_run = filter_modes(requested_modes, fn_name, supported_modes_for_fn)
        if not modes_to_run:
            warning(f"No valid modes to run for function '{fn_name}' based on request and function support. Skipping this function.")
            continue
        
        # create a function config for each backend and dtype combination
        for backend in supported_backends:
            for dtype in supported_dtypes:
                for mode in modes_to_run:
                    for env_config in supported_env_configs[backend]:
                        func_config = FunctionConfig(fn_name, mode, dtype, backend, env_config)
                        all_function_configs.append(func_config)
                        
                        # Generate inputs for this function configuration
                        fn_inputs = {}
                        for input_config in input_configs:
                            fn_inputs[input_config] = generate_fn_inputs(fn_name, *input_config, dtype, device)
                        
                        all_input_configs[func_config] = fn_inputs

    return all_function_configs, all_input_configs

    Parses command-line arguments and returns function configs and input configs.
    
flash_attn/flash_attn_triton_amd/bench.py