def load_flash_attn_module(backend: Literal["triton", "ck"], env_configs: Dict = {}, verbose = False):
    """
    Load the flash_attn module with the specified backend configuration
    """

    # remove any existing env variables first
    for key in ENV_FLAGS:
        if key in os.environ:
            del os.environ[key]

    # set environment variable for the desired backend
    if backend == "triton":
        os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE"
        os.environ["FLASH_ATTENTION_TRITON_AMD_AUTOTUNE"] = "0"
        os.environ["FLASH_ATTENTION_TRITON_AMD_DEBUG"] = "0"
    elif backend == "ck":
        os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "FALSE"
    else:
        raise ValueError(f"Unknown backend {backend}")
    
    # add custom env configs
    add_env_configs(env_configs)
    
    if verbose:
        print(f"Loading flash_attn module with {backend} backend.")
    
    # Remove any existing flash_attn modules from sys.modules
    for module_name in list(sys.modules.keys()):
        if module_name.startswith('flash_attn'):
            del sys.modules[module_name]
    
    # Clear CUDA cache
    torch.cuda.empty_cache()
    
    # Import and return the module
    import flash_attn
    
    return flash_attn

    Load the flash_attn module with the specified backend configuration
    
flash_attn/flash_attn_triton_amd/bench.py