def grouped_fmha_style_b2b_bmm_gen_function_call(func_attrs, indent="  "):
    """the function for generating a function call for attention"""
    assert len(func_attrs["outputs"]) == 1
    assert len(func_attrs["inputs"]) in (3, 4)

    output_name = func_attrs["outputs"][0]._attrs["name"]
    q_name = func_attrs["inputs"][0]._attrs["name"]
    k_name = func_attrs["inputs"][1]._attrs["name"]
    v_name = func_attrs["inputs"][2]._attrs["name"]

    bias_name = "nullptr"
    if len(func_attrs["inputs"]) == 4:
        bias_name = func_attrs["inputs"][3]._attrs["name"]

    q_shape = func_attrs["inputs"][0]._attrs["shape"]
    jagged_intvar = q_shape[0]
    batch_size_str = jagged_intvar.batch_dim()._attrs["name"]
    if len(jagged_intvar.jagged_dims()) != 1:
        raise RuntimeError(
            "Only support 1 jagged dim in grouped_fmha_style_b2b_bmm for now! "
            f"Current jagged intvar: {jagged_intvar}"
        )
    max_seq_length_dim = jagged_intvar.jagged_dims()[0].max_value()
    max_seq_length_str = (
        str(max_seq_length_dim.value())
        if isinstance(max_seq_length_dim, IntImm)
        else max_seq_length_dim._attrs["name"]
    )
    num_heads_str = q_shape[1]._attrs["name"]
    offset = f"{jagged_intvar.offsets_var_name()}.data[0]"

    return FUNC_CALL_TEMPLATE.render(
        func_name=func_attrs["name"],
        output=output_name,
        query=q_name,
        key=k_name,
        value=v_name,
        bias=bias_name,
        accum_ptr="global_workspace_",
        batch_size=batch_size_str,
        seq_length=max_seq_length_str,
        seq_length_kv=max_seq_length_str,
        num_heads=num_heads_str,
        offset=offset,
        indent=indent,
    )
the function for generating a function call for attention
python/aitemplate/backend/cuda/b2b_bmm/grouped_fmha_style_b2b_bmm.py