From cb2ca8591fb87dce58b5cd4e1ff3a4e36963efdf Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 23:44:53 +0000 Subject: [PATCH] fix: add @cute.jit to router compiled function --- dsv4/kernels/router/dense_router_decode_kernel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dsv4/kernels/router/dense_router_decode_kernel.py b/dsv4/kernels/router/dense_router_decode_kernel.py index f1ebcb94..c44bfc97 100644 --- 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