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

fix: use OperandMajorMode from nvgpu (not deprecated tcgen05) and mma_tiler_mn in router kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:39:50 +00:00
parent 1dbc57e2cd
commit 157f1c5258
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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