Repository: huggingface/transformers
Pull request: Cleaning hub kernels 

# What does this PR do?

Cleaning some comments, and renaming all instances of `load_and_register_kernel` to `load_and_register_attn_kernel` since it's used only for attention kernels

Changed files:
filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -27,7 +27,7 @@
 
 from ...configuration_utils import PreTrainedConfig
 from ...generation.configuration_utils import GenerationConfig
-from ...integrations.hub_kernels import load_and_register_kernel
+from ...integrations.hub_kernels import load_and_register_attn_kernel
 from ...utils.logging import logging
 from ...utils.metrics import ContinuousBatchProcessorMetrics, attach_tracer, traced
 from .cache import PagedAttentionCache
@@ -616,7 +616,7 @@ def __init__(
             if attn_implementation not in ALL_ATTENTION_FUNCTIONS._global_mapping:  # when its a kernel
                 from ...integrations.flash_paged import paged_attention_forward
 
-                load_and_register_kernel(attn_implementation, paged_attention_forward)
+                load_and_register_attn_kernel(attn_implementation, paged_attention_forward)
 
             model.config._attn_implementation = attn_implementation
         self.model = model.eval()

filename: src/transformers/integrations/hub_kernels.py
patch: @@ -42,7 +42,6 @@
         },
         "Llama4TextMoe": {
             "cuda": LayerRepository(
-                # Move to kernels-community/moe once we release.
                 repo_id="kernels-community/moe",
                 layer_name="Llama4TextMoe",
             )
@@ -51,13 +50,11 @@
             "cuda": LayerRepository(
                 repo_id="kernels-community/liger_kernels",
                 layer_name="LigerRMSNorm",
-                # revision="pure-layer-test",
             ),
             "rocm": {
                 Mode.INFERENCE: LayerRepository(
                     repo_id="kernels-community/liger_kernels",
                     layer_name="LigerRMSNorm",
-                    # revision="pure-layer-test",
                 )
             },
         },
@@ -169,7 +166,7 @@ def is_kernel(attn_implementation: Optional[str]) -> bool:
     )
 
 
-def load_and_register_kernel(attn_implementation: str, attention_wrapper: Optional[Callable] = None) -> None:
+def load_and_register_attn_kernel(attn_implementation: str, attention_wrapper: Optional[Callable] = None) -> None:
     """
     Load and register the kernel associated to `attn_implementation`.
 

filename: src/transformers/modeling_flash_attention_utils.py
patch: @@ -129,7 +129,7 @@ def lazy_import_flash_attention(implementation: Optional[str], force_import: Opt
     Lazily import flash attention and return the respective functions + flags.
 
     NOTE: For fullgraph, this needs to be called before compile, while no fullgraph can
-    work without preloading. See `load_and_register_kernel` in `integrations.hub_kernels`.
+    work without preloading. See `load_and_register_attn_kernel` in `integrations.hub_kernels`.
     """
     global _flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn
     if force_import or any(k is None for k in [_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn]):

filename: src/transformers/modeling_utils.py
patch: @@ -56,7 +56,7 @@
 from .integrations.flash_attention import flash_attention_forward
 from .integrations.flash_paged import paged_attention_forward
 from .integrations.flex_attention import flex_attention_forward
-from .integrations.hub_kernels import is_kernel, load_and_register_kernel
+from .integrations.hub_kernels import is_kernel, load_and_register_attn_kernel
 from .integrations.sdpa_attention import sdpa_attention_forward
 from .integrations.sdpa_paged import sdpa_attention_paged_forward
 from .integrations.tensor_parallel import (
@@ -2553,7 +2553,7 @@ def _check_and_adjust_attn_implementation(
 
         if is_kernel(applicable_attn_implementation):
             try:
-                load_and_register_kernel(applicable_attn_implementation)
+                load_and_register_attn_kernel(applicable_attn_implementation)
                 # log that we used kernel fallback if successful
                 if attn_implementation.startswith("flash_attention"):
                     logger.warning_once(