fix: add @cute.jit to router compiled function

2026-05-31 23:44:53 +00:00
parent d5d2b7b4b8
commit cb2ca8591f
1 changed files with 1 additions and 0 deletions
--- a/dsv4/kernels/router/dense_router_decode_kernel.py
+++ b/dsv4/kernels/router/dense_router_decode_kernel.py
@@ -109,6 +109,7 @@ class DenseRouterDecodeKernel:

        # All MLIR-dependent setup (tiled_mma, TMA atoms, CuTe tensor conversion)
        # 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):
            self._setup_attributes()
            tiled_mma = self._tiled_mma