def classic_b2b_bmm_gen_function(func_attrs: Dict[str, Any]) -> str:
    """the function for generating attention kernel"""
    q, k, v, bias = func_attrs["inputs"]
    seq_len_dim = 1
    n0 = k._attrs["shape"][seq_len_dim]
    n1 = v._attrs["shape"][-1]
    if not isinstance(n0, IntImm) or not isinstance(n1, IntImm):
        raise RuntimeError(
            f"n0 and n1 must be static dims. {func_attrs['name']=}, {n0=}, {n1=}"
        )

    supported_types = ("float16", "bfloat16")
    input_type = func_attrs["inputs"][0]._attrs["dtype"]
    output_type = func_attrs["outputs"][0]._attrs["dtype"]
    if input_type not in supported_types or output_type not in supported_types:
        raise NotImplementedError(
            f"{supported_types=} for inputs and output "
            f"but got {input_type=} and {output_type=}."
        )

    backend_spec = CUDASpec()
    elem_input_type = backend_spec.dtype_to_lib_type(input_type)
    elem_output_type = backend_spec.dtype_to_lib_type(output_type)

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

    import cutlass_lib

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

    bias_shape = bias._attrs["shape"]
    bias_broadcast = [s == IntImm(1) for s in bias_shape]
    if len(bias_broadcast) == 3:
        # single head case: Add num heads dimension of size 1
        bias_broadcast = [bias_broadcast[0], True, bias_broadcast[1], bias_broadcast[2]]
    assert (
        len(bias_broadcast) == 4
    ), f"Bias shape should be of length 4, got {len(bias_broadcast)=}"

    # Calculate stride expressions for bias tensor
    # Last dimension of bias has implicit stride of 1,
    # so cannot be broadcasted over
    bias_stride_n = "problem_size_0.n()"
    bias_shape_expr = [bias_stride_n]

    # build stride expressions
    if not bias_broadcast[-2]:
        bias_shape_expr.append("problem_size_0.m()")
    bias_stride_mn = "*".join(bias_shape_expr)
    if not bias_broadcast[-3]:
        bias_shape_expr.append("num_heads")
    bias_stride_hmn = "*".join(bias_shape_expr)  # batch stride

    # Strides for broadcasted dimensions are zero
    if bias_broadcast[0]:  # query sequence len stride
        bias_stride_hmn = "0"
    if bias_broadcast[1]:  # head stride
        bias_stride_mn = "0"
    if bias_broadcast[2]:  # query sequence length stride
        bias_stride_n = "0"

    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,
        n0=str(n0.value()),
        n1=str(n1.value()),
        has_causal=(
            "true" if func_attrs["causal_type"] != CausalType.NO_CAUSAL else "false"
        ),
        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"
        ),
        epilogue_math=epilogue_math,
        bias_stride_n=bias_stride_n,
        bias_stride_mn=bias_stride_mn,
        bias_stride_hmn=bias_stride_hmn,
    )
the function for generating attention kernel
python/aitemplate/backend/cuda/b2b_bmm/classic_b2b_bmm.py