From 824d054ad74ceaa70829295be320cf3a0910324a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 23:47:33 +0000 Subject: [PATCH] fix: inside cute.compile args are already CuTe tensors, no conversion needed --- dsv4/kernels/router/dense_router_decode_kernel.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/dsv4/kernels/router/dense_router_decode_kernel.py b/dsv4/kernels/router/dense_router_decode_kernel.py index a68e3e32..717615da 100644 --- a/dsv4/kernels/router/dense_router_decode_kernel.py +++ b/dsv4/kernels/router/dense_router_decode_kernel.py @@ -114,11 +114,12 @@ class DenseRouterDecodeKernel: self._setup_attributes() tiled_mma = self._tiled_mma - X_cu = cutlass_torch.from_dlpack(X).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(X)) - W_cu = cutlass_torch.from_dlpack(W_gate).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(W_gate)) - e_bias_cu = cutlass_torch.from_dlpack(e_bias).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(e_bias)) - out_w_cu = cutlass_torch.from_dlpack(out_w).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(out_w)) - out_ids_cu = cutlass_torch.from_dlpack(out_ids).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(out_ids)) + # Inside cute.compile, arguments are already CuTe tensors + X_cu = X + W_cu = W_gate + e_bias_cu = e_bias + out_w_cu = out_w + out_ids_cu = out_ids a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)