Skip to content

Commit babb128

Browse files
Elliot Gorokhovskyfacebook-github-bot
authored andcommitted
PR#4179 (#2435)
Summary: X-link: facebookresearch/FBGEMM#124 Pull Request resolved: #2435 X-link: pytorch/FBGEMM#3027 This PR is a dependency of the grid_constant PR. The API for TMA descriptor fill methods was changed, so I fixed up all usages in fbcode. triton-lang/triton#4179 Reviewed By: minjang Differential Revision: D61729239 fbshipit-source-id: 8ce25b7c230c3f4ad960f76aa0dd29626c8ee4d2
1 parent 52103b5 commit babb128

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

torchbenchmark/util/kernels/triton_fused_attention.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size):
6262
else:
6363
desc_x = self.cuda_descriptors[name]
6464
buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
65-
self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, buf_x.numpy())
65+
self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, buf_x.data_ptr())
6666
desc_x.copy_(buf_x, non_blocking=True)
6767

6868

@@ -75,7 +75,7 @@ def fill_2d_tma_descriptor(self, name, ptr, dim1, dim0, block_dim1, block_dim0,
7575
else:
7676
desc_x = self.cuda_descriptors[name]
7777
buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
78-
self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.numpy())
78+
self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr())
7979
desc_x.copy_(buf_x, non_blocking=True)
8080

8181

0 commit comments

Comments
 (0)