From a81612305c4c63790d7824ec2d2bf87f87c15052 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Tue, 23 Sep 2025 16:21:00 -0400 Subject: [PATCH] fix? --- src/axolotl/kernels/moe/tt_mg_gemm/tma_autotuning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/kernels/moe/tt_mg_gemm/tma_autotuning.py b/src/axolotl/kernels/moe/tt_mg_gemm/tma_autotuning.py index 2105ba518..f77ecbfeb 100644 --- a/src/axolotl/kernels/moe/tt_mg_gemm/tma_autotuning.py +++ b/src/axolotl/kernels/moe/tt_mg_gemm/tma_autotuning.py @@ -77,8 +77,8 @@ class TmaDescriptorHelper: ) self.tma_size = tma_size - self.fill_1d_tma_descriptor_inner = driver.active.utils.fill_1d_tma_descriptor - self.fill_2d_tma_descriptor_inner = driver.active.utils.fill_2d_tma_descriptor + self.fill_1d_tma_descriptor_inner = driver.active.utils.fill_tma_descriptor + self.fill_2d_tma_descriptor_inner = driver.active.utils.fill_tma_descriptor self.descriptors: Dict[str, torch.Tensor] = {} def init_tma_descriptor(self, name: str) -> None: