Skip to content

Fix #669: Support XPU and other backends in benchmark #727

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AdityaKulshrestha wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Fix #669: Support XPU and other backends in benchmark #727

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
4aa655c
Fixed hardcoded CUDA device to support device backend
AdityaKulshrestha Jul 25, 2026
da9be31
Fixed comment
AdityaKulshrestha Jul 25, 2026
e6be71e
Undo local changes
AdityaKulshrestha Jul 25, 2026
f91243b
linting
AdityaKulshrestha Oct 26, 2026
9c68669
chore: update dependencies to latest versions
AdityaKulshrestha Oct 26, 2026
5427db8
Merge branch 'main' into xpu-bench
AdityaKulshrestha Oct 26, 2026
0671e80
chores
AdityaKulshrestha Oct 26, 2026
5dd7b28
trailing newline fix
AdityaKulshrestha Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 39 additions & 39 deletions kernels/src/kernels/benchmarks/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,10 @@ class FlashAttentionBenchmark(Benchmark):
# Workload: small (B=2, S=128, H=8, D=64)
def setup_small(self):
B, S, H, D = 2, 128, 8, 64
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_small(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
Expand All @@ -58,10 +58,10 @@ def verify_small(self) -> torch.Tensor:
# Workload: medium (B=4, S=512, H=16, D=64)
def setup_medium(self):
B, S, H, D = 4, 512, 16, 64
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_medium(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
Expand All @@ -72,10 +72,10 @@ def verify_medium(self) -> torch.Tensor:
# Workload: large (B=8, S=1024, H=32, D=128)
def setup_large(self):
B, S, H, D = 8, 1024, 32, 128
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_large(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
Expand All @@ -90,10 +90,10 @@ class FlashAttentionCausalBenchmark(Benchmark):
# Workload: small (B=2, S=128, H=8, D=64)
def setup_small(self):
B, S, H, D = 2, 128, 8, 64
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_small(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
Expand All @@ -104,10 +104,10 @@ def verify_small(self) -> torch.Tensor:
# Workload: medium (B=4, S=512, H=16, D=64)
def setup_medium(self):
B, S, H, D = 4, 512, 16, 64
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_medium(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
Expand All @@ -118,10 +118,10 @@ def verify_medium(self) -> torch.Tensor:
# Workload: large (B=8, S=1024, H=32, D=128)
def setup_large(self):
B, S, H, D = 8, 1024, 32, 128
self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(B, S, H, D, device=self.device, dtype=torch.float16)
self.out = torch.empty(B, S, H, D, device=self.device, dtype=torch.float16)

def benchmark_large(self):
self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
Expand All @@ -139,16 +139,16 @@ def setup_small(self):
# Pack sequences of lengths [32, 48, 64]
seqlens = [32, 48, 64]
total = sum(seqlens)
self.q = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.cu_seqlens = torch.tensor(
[0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
device="cuda",
device=self.device,
dtype=torch.int32,
)
self.max_seqlen = max(seqlens)
self.out = torch.empty(total, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(total, H, D, device=self.device, dtype=torch.float16)

def benchmark_small(self):
self.out = _extract_output(
Expand All @@ -171,16 +171,16 @@ def setup_medium(self):
H, D = 16, 64
seqlens = [128, 192, 256, 200, 150]
total = sum(seqlens)
self.q = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.cu_seqlens = torch.tensor(
[0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
device="cuda",
device=self.device,
dtype=torch.int32,
)
self.max_seqlen = max(seqlens)
self.out = torch.empty(total, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(total, H, D, device=self.device, dtype=torch.float16)

def benchmark_medium(self):
self.out = _extract_output(
Expand All @@ -203,16 +203,16 @@ def setup_large(self):
H, D = 32, 128
seqlens = [256, 384, 512, 448, 320, 480, 400, 512]
total = sum(seqlens)
self.q = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.k = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.v = torch.randn(total, H, D, device="cuda", dtype=torch.float16)
self.q = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.k = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.v = torch.randn(total, H, D, device=self.device, dtype=torch.float16)
self.cu_seqlens = torch.tensor(
[0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
device="cuda",
device=self.device,
dtype=torch.int32,
)
self.max_seqlen = max(seqlens)
self.out = torch.empty(total, H, D, device="cuda", dtype=torch.float16)
self.out = torch.empty(total, H, D, device=self.device, dtype=torch.float16)

def benchmark_large(self):
self.out = _extract_output(
Expand Down
24 changes: 12 additions & 12 deletions kernels/src/kernels/benchmarks/layer_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@ class RMSNormBenchmark(Benchmark):
# Workload: small (B=2, S=128, D=768)
def setup_small(self):
B, S, D = 2, 128, 768
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down Expand Up @@ -40,8 +40,8 @@ def verify_small(self) -> torch.Tensor:
# Workload: medium (B=4, S=512, D=2048)
def setup_medium(self):
B, S, D = 4, 512, 2048
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down Expand Up @@ -70,8 +70,8 @@ def verify_medium(self) -> torch.Tensor:
# Workload: large (B=8, S=1024, D=4096)
def setup_large(self):
B, S, D = 8, 1024, 4096
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down Expand Up @@ -105,8 +105,8 @@ class LayerNormBenchmark(Benchmark):
# Workload: small (B=2, S=128, D=768)
def setup_small(self):
B, S, D = 2, 128, 768
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down Expand Up @@ -134,8 +134,8 @@ def verify_small(self) -> torch.Tensor:
# Workload: medium (B=4, S=512, D=2048)
def setup_medium(self):
B, S, D = 4, 512, 2048
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down Expand Up @@ -163,8 +163,8 @@ def verify_medium(self) -> torch.Tensor:
# Workload: large (B=8, S=1024, D=4096)
def setup_large(self):
B, S, D = 8, 1024, 4096
self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
self.x = torch.randn(B, S, D, device=self.device, dtype=torch.float16)
self.weight = torch.ones(D, device=self.device, dtype=torch.float16)
self.out = torch.empty_like(self.x)
self.B, self.S, self.D = B, S, D

Expand Down