Repository: huggingface/transformers
Pull request: :facepalm: CB nit! 

# What does this PR do?
Fix a small typo

Changed files:
filename: examples/pytorch/continuous_batching_simple.py
patch: @@ -33,15 +33,18 @@
     parser.add_argument("--max-batch-tokens", "-b", type=int, default=None)
     parser.add_argument("--attn", type=str, default="kernels-community/flash-attn", help="Attention implementation")
     parser.add_argument("--samples", type=int, default=500)
+    parser.add_argument("--max-new-tokens", type=int, default=32)
+
     args = parser.parse_args()
 
     # Prepare model
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         attn_implementation=args.attn,
+        device_map="cuda",
         dtype=torch.bfloat16,
     )
-    model = model.cuda().eval()
+    model = model.eval()
 
     # Prepare tokenizer and dataset
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
@@ -52,7 +55,7 @@
 
     # Prepare generation config
     generation_config = GenerationConfig(
-        max_new_tokens=512,
+        max_new_tokens=args.max_new_tokens,
         use_cuda_graph=False,  # Not supported for simple version
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.pad_token_id,

filename: src/transformers/generation/continuous_batching/cache.py
patch: @@ -173,7 +173,7 @@ def __init__(
         page_size = self.head_dim * self.num_key_value_heads
 
         if "flash" in self.config._attn_implementation:
-            num_attention_masks = 0
+            num_attention_masks = 1  # only used to compute the default meme args
         else:
             # TODO: when we generalize to allow for block-attn, we can use `num_attention_masks=sum(set(group_types))`
             num_attention_masks = 2 if "sliding_attention" in group_types else 1

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -27,6 +27,7 @@
 
 from ...configuration_utils import PreTrainedConfig
 from ...generation.configuration_utils import GenerationConfig
+from ...integrations.hub_kernels import load_and_register_kernel
 from ...utils.logging import logging
 from ...utils.metrics import ContinuousBatchProcessorMetrics, attach_tracer, traced
 from .cache import PagedAttentionCache
@@ -598,21 +599,27 @@ def __init__(
         streaming: bool = True,
         slice_inputs: bool = True,
     ):
-        """Initialize the continuous batching manager.
+        """
+        Initialize the continuous batching manager.
 
         Args:
             model: The language model for generation
             generation_config: Configuration for generation parameters
             max_queue_size: Maximum size of the request queue (0 = unlimited)
             streaming: Whether to stream tokens as they are generated
         """
-        self.model = model.eval()
         if "paged|" not in model.config._attn_implementation:
-            from ...integrations.hub_kernels import load_and_register_kernel
+            attn_implementation = f"paged|{model.config._attn_implementation}"
+
+            from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 
-            attn_implementation = "paged|" + self.model.config._attn_implementation
-            load_and_register_kernel(attn_implementation)
-            model.set_attn_implementation(attn_implementation)
+            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)
+
+            model.config._attn_implementation = attn_implementation
+        self.model = model.eval()
         generation_config = model.generation_config if generation_config is None else generation_config
         self.generation_config = generation_config
         self.input_queue = queue.Queue(maxsize=max_queue_size)

filename: src/transformers/integrations/hub_kernels.py
patch: @@ -13,7 +13,7 @@
 # limitations under the License.
 import re
 from functools import partial
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from ..modeling_flash_attention_utils import lazy_import_flash_attention
 from .flash_attention import flash_attention_forward
@@ -168,27 +168,29 @@ def is_kernel(attn_implementation: Optional[str]) -> bool:
     )
 
 
-def load_and_register_kernel(attn_implementation: str) -> None:
-    """Load and register the kernel associated to `attn_implementation`."""
-    if not is_kernel(attn_implementation):
+def load_and_register_kernel(attn_implementation: str, attention_wrapper: Optional[Callable] = None) -> None:
+    """
+    Load and register the kernel associated to `attn_implementation`.
+
+    Args:
+        attn_implementation: A string, usually a kernel repo like "kernels-community/flash-mla".
+        attn_wrapper: a callable for the wrapper around the attention implementation. In `transformers` we
+            have a wrapper around the `flash_attn_var_len` call, and the same goes for `sdpa` and `eager`.
+            They just prepare the arguments properly. This is mostly used for continious batching, where we
+            want the `paged` wrapper, which calls the paged cache.
+    """
+    from ..masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+    from ..modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    actual_attn_name = attn_implementation.split("|")[1] if "|" in attn_implementation else attn_implementation
+    if not is_kernel(actual_attn_name):
         return
     if not _kernels_available:
         raise ImportError(
             "`kernels` is either not installed or uses an incompatible version. "
             "Please install the latest version with `pip install -U kernels`."
         )
 
-    # Need to be imported here as otherwise we have a circular import in `modeling_utils`
-    from ..masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
-    from ..modeling_utils import ALL_ATTENTION_FUNCTIONS
-
-    attention_wrapper = None
-    # FIXME: @ArthurZucker this is dirty, did not want to do a lof of extra work
-    actual_attn_name = attn_implementation
-    if "|" in attn_implementation:
-        attention_wrapper, actual_attn_name = attn_implementation.split("|")
-        # `transformers` has wrapper for sdpa, paged, flash, flex etc.
-        attention_wrapper = ALL_ATTENTION_FUNCTIONS.get(attention_wrapper)
     # Extract repo_id and kernel_name from the string
     if ":" in actual_attn_name:
         repo_id, kernel_name = actual_attn_name.split(":")

filename: src/transformers/modeling_utils.py
patch: @@ -2576,12 +2576,10 @@ def _check_and_adjust_attn_implementation(
             # preload flash attention here to allow compile with fullgraph
             if applicable_attn_implementation.startswith("flash_attention"):
                 lazy_import_flash_attention(applicable_attn_implementation, force_import=True)
-
         return applicable_attn_implementation
 
     def get_correct_attn_implementation(self, requested_attention: Optional[str], is_init_check: bool = False) -> str:
         applicable_attention = "sdpa" if requested_attention is None else requested_attention
-
         if applicable_attention not in ["eager"] + ALL_ATTENTION_FUNCTIONS.valid_keys():
             message = (
                 f'Specified `attn_implementation="{applicable_attention}"` is not supported. The only possible arguments are '