Repository: huggingface/transformers
Pull request: feat: extract rev in attn_implementation kernels via @

This PR adds the ability to specify kernel revisions via the `@` symbol in  in the `attn_implementation` in `AutoModelForCausalLM.from_pretrained`


### Example usage

```bash
uv run repro.py
```

```python
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "accelerate",
#     "torch==2.7.0",
#     "transformers",
#     "kernels>=0.9.0",
# ]
#
# [tool.uv.sources]
# transformers = { git = "https://github.com/drbh/transformers.git", branch = "allow-kernel-rev" }
# ///
import time
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig


torch.set_float32_matmul_precision("high")

model_id = "meta-llama/Llama-3.2-3b-Instruct"
model = (
    AutoModelForCausalLM.from_pretrained(
        model_id,
        # attn_implementation="kernels-community/flash-attn@main",
        # attn_implementation="kernels-community/flash-attn@56449c1aa267bd0f48a191f0e6979dedf9f2ec32", # (mains sha)
        attn_implementation="kernels-community/flash-attn@09eec95", # main-1 commit sha
        torch_dtype=torch.bfloat16,
    )
    .eval()
    .cuda()
)
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")

print("+ Model loaded successfully.")

prompt = "What is the capital of France?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
generation_config = GenerationConfig(
    temperature=0.1,
    top_p=0.95,
    top_k=50,
    num_beams=1,
    max_new_tokens=50,
    do_sample=True,
    seed=42,
)
start_time = time.time()
with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        generation_config=generation_config,
    )
end_time = time.time()
print(f"+ Generation time: {end_time - start_time:.2f} seconds")
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print("+ Inference completed successfully.")
```

Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -2721,7 +2721,7 @@ def _check_and_adjust_attn_implementation(
             None to sdpa (to potentially eager).
         """
         applicable_attn_implementation = "sdpa" if attn_implementation is None else attn_implementation
-        if re.match(r"^[^/:]+/[^/:]+:?[^/:]+$", applicable_attn_implementation):
+        if re.match(r"^[^/:]+/[^/:]+(?:@[^/:]+)?(?::[^/:]+)?$", applicable_attn_implementation):
             if not is_kernels_available():
                 raise ValueError("kernels is not installed. Please install it with `pip install kernels`.")
             attention_wrapper = None
@@ -2738,8 +2738,12 @@ def _check_and_adjust_attn_implementation(
                 repo_id = applicable_attn_implementation
                 kernel_name = None
             repo_id = repo_id.strip()
+            # extract the rev after the @ if it exists
+            repo_id, _, rev = repo_id.partition("@")
+            repo_id = repo_id.strip()
+            rev = rev.strip() if rev else None
             try:
-                kernel = get_kernel(repo_id)
+                kernel = get_kernel(repo_id, revision=rev)
                 if hasattr(kernel, "flash_attn_varlen_func"):
                     if attention_wrapper is None:
                         attention_wrapper = flash_attention_forward
@@ -4494,6 +4498,22 @@ def from_pretrained(
             attn_implementation (`str`, *optional*):
                 The attention implementation to use in the model (if relevant). Can be any of `"eager"` (manual implementation of the attention), `"sdpa"` (using [`F.scaled_dot_product_attention`](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html)), `"flash_attention_2"` (using [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)), or `"flash_attention_3"` (using [Dao-AILab/flash-attention/hopper](https://github.com/Dao-AILab/flash-attention/tree/main/hopper)). By default, if available, SDPA will be used for torch>=2.1.1. The default is otherwise the manual `"eager"` implementation.
 
+                Accept HF kernel references in the form:
+                  <namespace>/<repo_name>[@<revision>][:<kernel_name>]
+
+                - <namespace> and <repo_name> are any non-"/" and non-":" sequences.
+                - "@<revision>" is optional (branch, tag, or commit-ish), e.g. "@main", "@v1.2.0", "@abc123".
+                - ":<kernel_name>" is optional and selects a function inside the kernel repo.
+                - Both options can appear together and in this order only: @revision first, then :kernel_name.
+                - We intentionally allow a leading "<wrapper>|" prefix (e.g., "flash|...") because the code
+                  strips it before loading; '|' is not excluded in the character classes here.
+
+                Examples that match:
+                  "org/model"
+                  "org/model@main"
+                  "org/model:custom_kernel"
+                  "org/model@v1.2.3:custom_kernel"
+
             > Parameters for big model inference
 
             torch_dtype (`str` or `torch.dtype`, *optional*):