From 157f1c5258094b26bb549568f52f7b446c96c1d9 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 23:39:50 +0000 Subject: [PATCH] fix: use OperandMajorMode from nvgpu (not deprecated tcgen05) and mma_tiler_mn in router kernel --- dsv4/kernels/router/dense_router_decode_kernel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dsv4/kernels/router/dense_router_decode_kernel.py b/dsv4/kernels/router/dense_router_decode_kernel.py index 42216f50..97b39b61 100644 --- a/dsv4/kernels/router/dense_router_decode_kernel.py +++ b/dsv4/kernels/router/dense_router_decode_kernel.py @@ -25,7 +25,7 @@ import torch import cutlass import cutlass.cute as cute -from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode import cutlass.utils as utils import cutlass.pipeline as pipeline import cutlass.utils.blackwell_helpers as sm100_utils @@ -101,8 +101,8 @@ class DenseRouterDecodeKernel: self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) def run(self, X, W_gate, e_bias, out_w, out_ids, M, E, K, scaling, top_k, stream=None): - self.a_major_mode = tcgen05.OperandMajorMode.K - self.b_major_mode = tcgen05.OperandMajorMode.K + self.a_major_mode = OperandMajorMode.K + self.b_major_mode = OperandMajorMode.K self._setup_attributes() X_cu = cutlass_torch.to_cuTe_tensor(X, major_mode=self.a_major_mode)