def flash_attention_gen_function(func_attrs: Dict[str, Any]) -> str:
    """the function for generating attention kernel"""
    return FUNC_TEMPLATE.render(
        custom_kernel=ATT_KERNEL_TEMPLATE.render(
            s1=128 if func_attrs["seq_len"] == 128 else 256,
            s2=func_attrs["head_size"],
        ),
        func_signature=FUNC_SIGNATURE.render(func_name=func_attrs["name"]),
    )
the function for generating attention kernel
python/aitemplate/backend/cuda/attention/flash_attention.py