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

fix: add @cute.jit to router compiled function

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:44:53 +00:00
parent d5d2b7b4b8
commit cb2ca8591f
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -109,6 +109,7 @@ class DenseRouterDecodeKernel:
# All MLIR-dependent setup (tiled_mma, TMA atoms, CuTe tensor conversion) # All MLIR-dependent setup (tiled_mma, TMA atoms, CuTe tensor conversion)
# must happen inside cute.compile context. This matches the MoE kernel pattern. # must happen inside cute.compile context. This matches the MoE kernel pattern.
@cute.jit
def _compiled_fn(X, W_gate, e_bias, out_w, out_ids): def _compiled_fn(X, W_gate, e_bias, out_w, out_ids):
self._setup_attributes() self._setup_attributes()
tiled_mma = self._tiled_mma tiled_mma = self._tiled_mma