[XPU] Add flash_attn2 support for XPU
Author: YangKai0616
Pull request: yes

# What does this PR do?

XPU has now implemented the basic functionality of `flash_attn2` in [kernels-community/flash-attn2](https://huggingface.co/kernels-community/flash-attn2/tree/main/build). So this PR adds the function call to transformers. 

Main contributions:

1. Added `flash_attn2` support for **XPU**;
2. Enabled `flash_attn2` UTs on **XPU**.

Now we can use `attn_implementation="flash_attention_2"` or `attn_implementation="kernels-community/flash-attn2"` on **XPU** to invoke this feature.

**Note** (CI):

1. Some tests such as `XXX::test_flash_attn_2_inference_equivalence` and `XXX::test_flash_attn_2_equivalence`, etc. , will sometimes pass and sometimes fail due to randomness introduced by the input data and model initialization. I have observed this phenomenon on both **CUDA (A100)** and **XPU**.
2. For the test case `tests/models/kosmos2/test_modeling_kosmos2.py::Kosmos2ModelTest::test_eager_matches_fa2_generate`, etc., **CUDA** will trigger a `RuntimeError: cu_seqlens_q must have shape (batch_size + 1)`, while **XPU** will directly cause a `Aborted (core dumped)`. This appears to be caused by the robustness of the underlying data reading function, but the root cause is the test case issue.

Comments:
Once [huggingface/kernels-community/pull/59](https://github.com/huggingface/kernels-community/pull/59) has been merged and the binary files has been uploaded to [kernels-community/flash-attn2](https://huggingface.co/kernels-community/flash-attn2/tree/main/build), we can begin the review. Cc @yao-matrix .

@YangKai0616 , I think the latest kernel binaries have been properly, let's verify and kick off this PR, thx.

I compared the performance of **flash_attn2** and **sdpa** on XPU using the following code. The results show that flash_attn2 achieves an average speedup of **2.59x**.

```
import math
import time
import torch
import torch.nn.functional as F

try:
    from kernels import get_kernel
    flash_attn = get_kernel("kernels-community/flash-attn2")
    flash_attn_varlen_func = flash_attn.flash_attn_varlen_func

    FLASH_ATTN_AVAILABLE = True
except ImportError:
    FLASH_ATTN_AVAILABLE = False
    print("Warning: flash_attn_varlen_func not available")

speedups = []

def generate_random_padding_mask(max_seqlen, batch_size, device, mode="random"):
    assert mode in ["full", "random", "third"]
    if mode == "full":
        lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)
    elif mode == "random":
        lengths = torch.randint(
            max(1, max_seqlen - 20), max_seqlen + 1, (batch_size, 1), device=device
        )
    elif mode == "third":
        lengths = torch.randint(max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device)
    
    padding_mask = (
        torch.arange(max_seqlen, device=device).unsqueeze(0) < lengths
    )
    return padding_mask


def unpad_input(hidden_states, attention_mask):
    seqlen = attention_mask.shape[-1]
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    max_seqlen = attention_mask.sum(dim=-1).max().item()
    cu_seqlens = torch.cat([
        torch.tensor([0], dtype=torch.int32, device=attention_mask.device),
        attention_mask.sum(dim=-1).cumsum(dim=0).to(dtype=torch.int32)
    ])
    
    hidden_states_unpad = hidden_states.view(-1, *hidden_states.shape[2:])[indices]
    return hidden_states_unpad, indices, cu_seqlens, max_seqlen


def pad_input(hidden_states_unpad, indices, batch, seqlen):
    output = torch.zeros(
        batch * seqlen, *hidden_states_unpad.shape[1:],
        dtype=hidden_states_unpad.dtype,
        device=hidden_states_unpad.device
    )
    output[indices] = hidden_states_unpad
    return output.view(batch, seqlen, *hidden_states_unpad.shape[1:])


def sdpa_varlen_attention(
    q, k, v,
    query_padding_mask, key_padding_mask,
    dropout_p=0.0,
    softmax_scale=None,
    causal=False,
):
    batch_size, seqlen_q, nheads, d = q.shape
    seqlen_k = k.shape[1]
    nheads_k = k.shape[2]
    
    if softmax_scale is None:
        softmax_scale = 1.0 / math.sqrt(d)
    
    # Convert dimensions to SDPA expected format: (batch, nheads, seqlen, d)
    q = q.transpose(1, 2)  # (batch, nheads, seqlen_q, d)
    k = k.transpose(1, 2)  # (batch, nheads_k, seqlen_k, d)
    v = v.transpose(1, 2)  # (batch, nheads_k, seqlen_k, d)
    
    # For GQA/MQA, need to expand k and v
    if nheads != nheads_k:
        n_rep = nheads // nheads_k
        k = k.repeat_interleave(n_rep, dim=1)
        v = v.repeat_interleave(n_rep, dim=1)
    
    # Build complete attention mask (batch, 1, seqlen_q, seqlen_k)
    # Initialize to all True
    attn_mask = torch.ones(batch_size, 1, seqlen_q, seqlen_k, 
                           dtype=torch.bool, device=q.device)
    
    # Apply key padding mask: set padding positions to False
    if key_padding_mask is not None:
        # key_padding_mask: (batch, seqlen_k)
        # Expand to (batch, 1, 1, seqlen_k) and apply
        attn_mask = attn_mask & key_padding_mask.unsqueeze(1).unsqueeze(2)
    
    # Apply causal mask
    if causal:
        # Build causal mask separately for each batch
        # For varlen, query position i can only see key position j <= i
        causal_mask = torch.tril(
            torch.ones(seqlen_q, seqlen_k, dtype=torch.bool, device=q.device)
        )
        # Broadcast apply: (batch, 1, seqlen_q, seqlen_k) & (seqlen_q, seqlen_k)
        attn_mask = attn_mask & causal_mask.unsqueeze(0).unsqueeze(0)
    
    # Apply query padding mask: set entire rows corresponding to query padding positions to False
    if query_padding_mask is not None:
        # query_padding_mask: (batch, seqlen_q)
        # Expand to (batch, 1, seqlen_q, 1) and apply
        attn_mask = attn_mask & query_padding_mask.unsqueeze(1).unsqueeze(3)
    
    # Call SDPA
    # attn_mask is now (batch, 1, seqlen_q, seqlen_k)
    output = F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=attn_mask,
        dropout_p=dropout_p,
        scale=softmax_scale,
        is_causal=False  # We've already manually created the mask
    )
    
    # Convert back to (batch, seqlen_q, nheads, d)
    output = output.transpose(1, 2)
    
    # Zero out padding positions in output (double safety)
    if query_padding_mask is not None:
        output = output * query_padding_mask.unsqueeze(-1).unsqueeze(-1)
    
    return output


def benchmark_function(func, *args, num_warmup=10, num_runs=50, **kwargs):
    device = args[0].device
    
    # Warmup
    for _ in range(num_warmup):
        _ = func(*args, **kwargs)
    
    # Synchronize
    if device.type == 'cuda':
        torch.cuda.synchronize(device)
    elif device.type == 'xpu':
        torch.xpu.synchronize(device)
    
    # Actual testing
    start_time = time.perf_counter()
    for _ in range(num_runs):
        _ = func(*args, **kwargs)
    
    # Synchronize
    if device.type == 'cuda':
        torch.cuda.synchronize(device)
    elif device.type == 'xpu':
        torch.xpu.synchronize(device)
    
    end_time = time.perf_counter()
    
    avg_time = (end_time - start_time) / num_runs
    return avg_time


def test_flash_attn_varlen_perf(
    batch_size=8,
    seqlen_q=512,
    seqlen_k=512,
    nheads=12,
    nheads_k=12,
    d=64,
    causal=True,
    dtype=torch.float16,
    device='cuda',
    num_warmup=10,
    num_runs=50,
):
    print("=" * 80)
    print(f"Flash Attention Varlen Performance Test")
    print("=" * 80)
    print(f"Configuration:")
    print(f"  Device: {device}")
    print(f"  Data type: {dtype}")
    print(f"  Batch size: {batch_size}")
    print(f"  Seq len (Q): {seqlen_q}")
    print(f"  Seq len (K): {seqlen_k}")
    print(f"  Num heads (Q): {nheads}")
    print(f"  Num heads (K): {nheads_k}")
    print(f"  Head dim: {d}")
    print(f"  Causal: {causal}")
    print(f"  Warmup runs: {num_warmup}")
    print(f"  Benchmark runs: {num_runs}")
    print()
    
    # Set random seed
    torch.manual_seed(0)
    
    # Generate random data
    q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype)
    k = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype)
    v = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype)
    
    # Generate padding mask (use partial padding to simulate real scenario)
    query_padding_mask = generate_random_padding_mask(seqlen_q, batch_size, device, mode="random")
    key_padding_mask = query_padding_mask.clone()
    
    # Unpad input
    q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, query_padding_mask)
    k_unpad, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
    v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
    
    print(f"Actual data statistics:")
    print(f"  Total Q tokens: {q_unpad.shape[0]} (avg per batch: {q_unpad.shape[0] / batch_size:.1f})")
    print(f"  Total K tokens: {k_unpad.shape[0]} (avg per batch: {k_unpad.shape[0] / batch_size:.1f})")
    print(f"  Max Q seqlen: {max_seqlen_q}")
    print(f"  Max K seqlen: {max_seqlen_k}")
    print()
    
    # Test SDPA
    print("Testing PyTorch SDPA...")
    try:
        sdpa_time = benchmark_function(
            sdpa_varlen_attention,
            q, k, v,
            query_padding_mask, key_padding_mask,
            dropout_p=0.0,
            softmax_scale=None,
            causal=causal,
            num_warmup=num_warmup,
            num_runs=num_runs,
        )
        print(f"✓ SDPA average time: {sdpa_time * 1000:.3f} ms")
        sdpa_success = True
    except Exception as e:
        print(f"✗ SDPA failed: {e}")
        sdpa_success = False
        sdpa_time = None
    
    print()
    
    # Test Flash Attention Varlen
    if FLASH_ATTN_AVAILABLE:
        print("Testing Flash Attention Varlen...")
        try:
            flash_time = benchmark_function(
                flash_attn_varlen_func,
                q_unpad, k_unpad, v_unpad,
                cu_seqlens_q, cu_seqlens_k,
                max_seqlen_q, max_seqlen_k,
                dropout_p=0.0,
                softmax_scale=None,
                causal=causal,
                num_warmup=num_warmup,
                num_runs=num_runs,
            )
            print(f"✓ Flash Attention Varlen average time: {flash_time * 1000:.3f} ms")
            flash_success = True
        except Exception as e:
            print(f"✗ Flash Attention Varlen failed: {e}")
            flash_success = False
            flash_time = None
    else:
        print("✗ Flash Attention not available")
        flash_success = False
        flash_time = None
    
    print()
    print("=" * 80)
    print("Performance Comparison Summary:")
    print("=" * 80)
    
    if sdpa_success:
        print(f"SDPA:                     {sdpa_time * 1000:8.3f} ms")
    else:
        print(f"SDPA:                     Failed")
    
    if flash_success:
        print(f"Flash Attention Varlen:   {flash_time * 1000:8.3f} ms")
    else:
        print(f"Flash Attention Varlen:   Failed")
    
    if sdpa_success and flash_success:
        speedup = sdpa_time / flash_time
        print()
        if speedup > 1:
            speedups.append(speedup)
            print(f"Speedup: {speedup:.2f}x (Flash Attention is {(speedup - 1) * 100:.1f}% faster)")
        else:
            print(f"Speedup: {speedup:.2f}x (SDPA is {(1 / speedup - 1) * 100:.1f}% faster)")
    
    print("=" * 80)
    print()
    
    # Verify correctness
    if sdpa_success and flash_success:
        print("Verifying numerical correctness...")
        with torch.no_grad():
            out_sdpa = sdpa_varlen_attention(
                q, k, v,
                query_padding_mask, key_padding_mask,
                causal=causal,
            )
            out_flash = flash_attn_varlen_func(
                q_unpad, k_unpad, v_unpad,
                cu_seqlens_q, cu_seqlens_k,
                max_seqlen_q, max_seqlen_k,
                causal=causal,
            )

            out_flash_padded = pad_input(out_flash, indices_q, batch_size, seqlen_q)
        
        max_diff = (out_sdpa - out_flash_padded).abs().max().item()
        mean_diff = (out_sdpa - out_flash_padded).abs().mean().item()
        
        print(f"Max output difference: {max_diff:.6f}")
        print(f"Mean output difference: {mean_diff:.6f}")
        
        if max_diff < 0.01:
            print("✓ Numerical verification passed")
        else:
            print("✗ Large numerical difference, please check")
        print()


def run_comprehensive_tests():
    if torch.cuda.is_available():
        device = 'cuda'
        print(f"Using CUDA device: {torch.cuda.get_device_name()}")
    elif torch.xpu.is_available():
        device = 'xpu'
        print(f"Using XPU device")
    else:
        print("No GPU device available")
        return
    
    print()
    
    # Test configuration list
    test_configs = [
        # (batch_size, seqlen_q, seqlen_k, nheads, nheads_k, d, causal, dtype, description)
        # Small sequence length + small head dimension
        (4, 147, 147, 4, 4, 32, True, torch.bfloat16, "seq(147,147) d=32 MHA"),
        (4, 113, 113, 4, 1, 64, True, torch.bfloat16, "seq(113,113) d=64 MQA"),
        
        # Small sequence length + medium head dimension
        (4, 128, 128, 4, 2, 96, True, torch.bfloat16, "seq(128,128) d=96 GQA"),
        (4, 113, 113, 4, 4, 128, True, torch.bfloat16, "seq(113,113) d=128 MHA"),
        
        # Medium sequence length + small head dimension
        (4, 256, 256, 4, 1, 59, True, torch.bfloat16, "seq(256,256) d=59 MQA"),
        (4, 512, 512, 4, 4, 64, True, torch.bfloat16, "seq(512,512) d=64 MHA"),
        
        # Medium sequence length + medium head dimension
        (4, 512, 512, 4, 2, 80, True, torch.bfloat16, "seq(512,512) d=80 GQA"),
        (4, 256, 256, 4, 4, 128, True, torch.bfloat16, "seq(256,256) d=128 MHA"),
        
        # Medium sequence length + large head dimension
        (4, 512, 512, 4, 1, 160, True, torch.bfloat16, "seq(512,512) d=160 MQA"),
        (4, 256, 256, 4, 2, 192, True, torch.bfloat16, "seq(256,256) d=192 GQA"),
        
        # Large sequence length + small/medium head dimension
        (4, 1024, 1024, 4, 4, 64, True, torch.bfloat16, "seq(1024,1024) d=64 MHA"),
        
        # Large sequence length + large head dimension
        (4, 1024, 1024, 4, 2, 224, True, torch.bfloat16, "seq(1024,1024) d=224 GQA"),
        (4, 2048, 2048, 4, 4, 128, True, torch.bfloat16, "seq(2048,2048) d=128 MHA"),
        
        # Extra large head dimension test
        (4, 1024, 1024, 4, 1, 256, True, torch.bfloat16, "seq(1024,1024) d=256 MQA"),
    ]
    
    
    for config in test_configs:
        batch_size, seqlen_q, seqlen_k, nheads, nheads_k, d, causal, dtype, desc = config
        print(f"\n{'='*80}")
        print(f"Test: {desc}")
        print(f"{'='*80}\n")
        
        try:
            test_flash_attn_varlen_perf(
                batch_size=batch_size,
                seqlen_q=seqlen_q,
                seqlen_k=seqlen_k,
                nheads=nheads,
                nheads_k=nheads_k,
                d=d,
                causal=causal,
                dtype=dtype,
                device=device,
                num_warmup=15,
                num_runs=50,
            )
        except Exception as e:
            print(f"Test failed: {e}")
            import traceback
            traceback.print_exc()
        
        print()


if __name__ == "__main__":
    if torch.cuda.is_available():
        device = 'cuda'
    elif torch.xpu.is_available():
        device = 'xpu'
    else:
        print("No GPU device available")
        exit(1)

    # Run comprehensive tests
    run_comprehensive_tests()
    print(f"\nAll tests completed, average speedup: {sum(speedups)/len(speedups):.2f}x")
```

Hi @vasqu , we can start the review now. Thanks.

> Very cool to see this tho! Are there also plans for fa3 kernels?

Yes, our team is developing FA3.

**CI:**

- The failed example in `tests_processors` seems unrelated to this PR.
- The failed example `tests/models/resnet/test_modeling_resnet.py::ResNetModelTest::test_can_load_ignoring_mismatched_shapes` in `tests_torch`, I cannot reproduce it locally on either CUDA (A100) or XPU — both pass normally. Test environment: torch 2.9.0 and torch 2.9.1+xpu.

Hey @YangKai0616, 
you're right about the tests, they are unrelated to this PR / are resolved on main. Checking in with a team member to fix the failing test, but nothing to be done on your side. I will merge when we resolved the issue :hugs: 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_41956). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: bamba, diffllama, ernie4_5_moe, esm, glm4, gpt2, jamba, jetmoe, kosmos2_5, m2m_100, ministral, mixtral, modernbert, seed_oss, zamba, zamba2