def flash_attention_gen_function_call(func_attrs, indent="  "):
    """the function for generating a function call for attention"""
    output_name = ""
    assert len(func_attrs["outputs"]) == 1
    assert len(func_attrs["inputs"]) == 2

    output_name = func_attrs["outputs"][0]._attrs["name"]

    qkv_name = func_attrs["inputs"][0]._attrs["name"]

    seqlens_name = FUNC_CALL_INT32_PARAM_TEMPLATE.render(
        name=func_attrs["inputs"][1]._attrs["name"]
    )

    x = func_attrs["inputs"][0]

    batch_size = func_attrs["batch_size"]
    seq_len = func_attrs["seq_len"]

    num_heads = x._attrs["shape"][2]._attrs["values"][0]
    head_size = x._attrs["shape"][3]._attrs["values"][0]
    p_dropout = func_attrs["dropout"]
    is_causal = func_attrs["causal"]
    softmax_scale = head_size ** (-0.5)

    return FUNC_CALL_TEMPLATE.render(
        func_name=func_attrs["name"],
        output=output_name,
        qkv=qkv_name,
        cu_seqlens=seqlens_name,
        softmax_lse="reinterpret_cast<float*>(global_workspace_)",
        o_tmp="reinterpret_cast<float*>(global_workspace_ + {} * sizeof(float))".format(
            batch_size * num_heads * seq_len
        ),
        batch_size=batch_size,
        seq_len=seq_len,
        num_heads=num_heads,
        head_size=head_size,
        p_dropout=p_dropout,
        softmax_scale=softmax_scale,
        is_causal="true" if is_causal else "false",
        loop="true" if seq_len > 256 else "false",
        indent=indent,
    )
the function for generating a function call for attention
python/aitemplate/backend/cuda/attention/flash_attention.py