From 8e02d677bc8efa2134bfa14d18bd3581b51590ae Mon Sep 17 00:00:00 2001 From: Taras Sereda Date: Mon, 29 Sep 2025 16:50:40 -0700 Subject: [PATCH 1/2] remote device from sdpa tensors --- KernelBench/level1/97_ScaledDotProductAttention.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/KernelBench/level1/97_ScaledDotProductAttention.py b/KernelBench/level1/97_ScaledDotProductAttention.py index feff1490..bcaf939e 100644 --- a/KernelBench/level1/97_ScaledDotProductAttention.py +++ b/KernelBench/level1/97_ScaledDotProductAttention.py @@ -15,9 +15,9 @@ def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Te embedding_dimension = 1024 def get_inputs(): - Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, device='cuda', dtype=torch.float16) - K = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, device='cuda', dtype=torch.float16) - V = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, device='cuda', dtype=torch.float16) + Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) + K = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) + V = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) return [Q, K, V] def get_init_inputs(): From 52fa3018af9c4307edf262e5f7614de097db05bf Mon Sep 17 00:00:00 2001 From: Taras Sereda Date: Mon, 29 Sep 2025 16:54:09 -0700 Subject: [PATCH 2/2] dtype as well --- KernelBench/level1/97_ScaledDotProductAttention.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/KernelBench/level1/97_ScaledDotProductAttention.py b/KernelBench/level1/97_ScaledDotProductAttention.py index bcaf939e..ba4bd02b 100644 --- a/KernelBench/level1/97_ScaledDotProductAttention.py +++ b/KernelBench/level1/97_ScaledDotProductAttention.py @@ -15,9 +15,9 @@ def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Te embedding_dimension = 1024 def get_inputs(): - Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) - K = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) - V = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension, dtype=torch.float16) + Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension) + K = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension) + V = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension) return [Q, K, V] def get_init_inputs():