[AdityaKulshrestha]

Closes #669

What does this PR do?

• Fixes the hardcoded benchmarks for CUDA.
• Makes the benchmark system backend-agnostic by using the dynamically detected device.

Changes

• Added backend parameter to get_kernel() and get_local_kernel() to get the supported backend.
• Replaced all device="cuda" with device=self.device.

Testing

• Tested with flash-attn benchmarking on XPU - kernels benchmark kernels-community/flash-attn2 --version 1
• Output

Downloading kernels-community/flash-attn2@v1...
Running benchmark.py...

  GPU      N/A
  CPU      x86_64
  OS       Linux 7.0.0-27-generic
  PyTorch  2.12.0+xpu

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
  Running FlashAttn on xpu (backend: xpu)
  Running FlashAttnCausal on xpu (backend: xpu)
  Running FlashAttnVarlen on xpu (backend: xpu)

┌─────────────────┬────────────┬─────┬───────────┬────────────┬───────────┬───────────┬───────────┬───────────┬────────────┬───────────┬─────────┐
│ Benchmark       │ Workload   │   N │ Speedup   │   Mean(ms) │   Std(ms) │   Min(ms) │   Max(ms) │   IQR(ms) │   Outliers │   Ref(ms) │ Match   │
├─────────────────┼────────────┼─────┼───────────┼────────────┼───────────┼───────────┼───────────┼───────────┼────────────┼───────────┼─────────┤
│ FlashAttn       │ large      │ 100 │ 11.81x    │     1.6025 │    0.1022 │    1.5189 │    2.2015 │    0.0649 │          8 │   18.9221 │ ✓       │
│ FlashAttn       │ medium     │ 100 │ 11.37x    │     0.1015 │    0.0068 │    0.0928 │    0.1424 │    0.0042 │          6 │    1.1543 │ ✓       │
│ FlashAttn       │ small      │ 100 │ 2.65x     │     0.0723 │    0.0033 │    0.0689 │    0.0851 │    0.0024 │         10 │    0.1917 │ ✓       │
│ FlashAttnCausal │ large      │ 100 │ 15.19x    │     1.2892 │    0.0455 │    1.2097 │    1.3777 │    0.0867 │          0 │   19.5821 │ ✓       │
│ FlashAttnCausal │ medium     │ 100 │ 11.69x    │     0.1046 │    0.0056 │    0.0973 │    0.1295 │    0.0049 │          6 │    1.2228 │ ✓       │
│ FlashAttnCausal │ small      │ 100 │ 3.34x     │     0.0722 │    0.0051 │    0.0677 │    0.1105 │    0.0023 │         13 │    0.2415 │ ✓       │
│ FlashAttnVarlen │ large      │ 100 │ 13.76x    │     0.3795 │    0.0225 │    0.3604 │    0.5332 │    0.0138 │          5 │    5.2222 │ ✓       │
│ FlashAttnVarlen │ medium     │ 100 │ 19.14x    │     0.0771 │    0.003  │    0.0731 │    0.0965 │    0.0018 │          9 │    1.4756 │ ✓       │
│ FlashAttnVarlen │ small      │ 100 │ 11.42x    │     0.0775 │    0.0028 │    0.0741 │    0.0869 │    0.0026 │         12 │    0.8851 │ ✓       │
└─────────────────┴────────────┴─────┴───────────┴────────────┴───────────┴───────────┴───────────┴───────────┴────────────┴───────────┴─────────┘

  large: 11.81x faster (95% CI: 1.5825-1.6225ms vs ref 18.9221ms) ✓ significant
  medium: 11.37x faster (95% CI: 0.1002-0.1028ms vs ref 1.1543ms) ✓ significant
  small: 2.65x faster (95% CI: 0.0717-0.0729ms vs ref 0.1917ms) ✓ significant
  large: 15.19x faster (95% CI: 1.2803-1.2981ms vs ref 19.5821ms) ✓ significant
  medium: 11.69x faster (95% CI: 0.1035-0.1057ms vs ref 1.2228ms) ✓ significant
  small: 3.34x faster (95% CI: 0.0712-0.0732ms vs ref 0.2415ms) ✓ significant
  large: 13.76x faster (95% CI: 0.3751-0.3839ms vs ref 5.2222ms) ✓ significant
  medium: 19.14x faster (95% CI: 0.0765-0.0777ms vs ref 1.4756ms) ✓ significant
  small: 11.42x faster (95% CI: 0.0770-0.0780ms vs ref 0.8851ms) ✓ significant

Kernel: f12afc9  Benchmark: 685a205

Checklist

• This PR is linked to an issue that was discussed and approved
• I have tested these changes locally
• New/changed functionality has test coverage
• LLM disclosure:
  • I did not use an LLM to create this PR.
  • I used and LLM for assistance while creating this PR.
  • This PR was mostly or completely generated by an LLM.

[github-actions]

Hi @AdityaKulshrestha, thanks for your interest in contributing!

This project requires that pull request authors are vouched, and you are not in the list of vouched users.

This PR will be closed automatically. See https://github.com/huggingface/kernels/blob/main/CONTRIBUTING.md for more details.

[AdityaKulshrestha]

Hi @sayakpaul please review.

[AdityaKulshrestha]

Hi @sayakpaul @danieldk
Please review.

[danieldk]

Thanks for fixing this!

[AdityaKulshrestha]

@danieldk Made the suggested changes. Could you please take another look?