biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: inside cute.compile args are already CuTe tensors, no conversion needed

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:47:33 +00:00
parent 6375e54396
commit 824d054ad7
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
dsv4/kernels/router/dense_router_decode_kernel.py
View File

@@ -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)