def classic_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"]) == 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 = func_attrs["inputs"][3]._attrs["name"]

    q_shape = func_attrs["inputs"][0]._attrs["shape"]

    batch_size = q_shape[0]._attrs["name"]
    seq_len_dim = 1
    head_dim = -2
    m0 = q_shape[seq_len_dim]._attrs["name"]

    if len(q_shape) == 3:
        # single head case
        k0 = q_shape[2]._attrs["name"]
        num_heads = "1"
    elif len(q_shape) == 4:
        k0 = q_shape[3]._attrs["name"]
        num_heads = q_shape[head_dim]._attrs["name"]
    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,
        batch_size=batch_size,
        num_heads=num_heads,
        m0=m0,
        k0=k0,
        indent=indent,
    )
the function for generating a function call for attention
python/aitemplate/backend/cuda/b2b_bmm/classic_b2b_bmm.py