def fmha_style_b2b_bmm_gen_function(func_attrs: Dict[str, Any]) -> str:
    """the function for generating attention kernel"""
    q, k, v = func_attrs["inputs"][0:3]

    bias_broadcast = [False] * 4
    if len(func_attrs["inputs"]) > 3:
        bias = func_attrs["inputs"][3]
        bias_broadcast = [var == IntImm(1) for var in bias.shape()]

    k0 = k._attrs["shape"][3]
    n1 = v._attrs["shape"][3]
    if not isinstance(k0, IntImm) or not isinstance(n1, IntImm):
        raise RuntimeError(
            f"k0 and n1 must be static dims. {func_attrs['name']=}, {k0=}, {n1=}"
        )
    backend_spec = CUDASpec()
    elem_input_type = backend_spec.dtype_to_lib_type(
        func_attrs["inputs"][0]._attrs["dtype"]
    )
    elem_output_type = backend_spec.dtype_to_lib_type(
        func_attrs["outputs"][0]._attrs["dtype"]
    )

    if (
        "use_fp16_acc" in Target.current()._kwargs
        and Target.current()._kwargs["use_fp16_acc"]
        and elem_input_type == "cutlass::half_t"
    ):
        elem_accum_type = "cutlass::half_t"
    else:
        elem_accum_type = "float"

    import cutlass_lib

    activation_functor = cutlass_lib.library.EpilogueMathTag[
        cutlass_lib.library.EpilogueMathName[func_attrs["epilogue_math_name"]]
    ]

    return FUNC_TEMPLATE.render(
        func_name=func_attrs["name"],
        func_signature=FUNC_SIGNATURE.render(func_name=func_attrs["name"]),
        elem_input_type=elem_input_type,
        elem_output_type=elem_output_type,
        elem_accum_type=elem_accum_type,
        offset_t="int64_t",
        head_dim=str(k0.value()),
        head_dim_value=str(n1.value()),
        causal_type=causal_type_to_kernel_str(func_attrs["causal_type"]),
        alpha0=str(func_attrs["alpha0"]),
        alpha1=str(func_attrs["alpha1"]),
        alpha1_divide_by_seq_len=(
            "true" if func_attrs["alpha1_divide_by_seq_len"] else "false"
        ),
        activation_functor=activation_functor,
        bias_broadcast=bias_broadcast,
        offset_ptr="nullptr",
    )
the function for generating attention kernel
python/aitemplate/backend/cuda/b2b_bmm/fmha_style_b2b_bmm.py