From b536f9919299133b11fc671adde11559ece23f96 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 09:32:05 +0000 Subject: [PATCH] CRITICAL FIX: move ALL CuTe DSL setup inside @cute.jit context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The root cause of ALL the MLIR crashes: _create_tiled_mma and _setup_attributes call cute.make_tiled_mma, sm100_utils.make_smem_layout_a, etc. These are MLIR operations that REQUIRE an active MLIR context. Previously they ran in run() OUTSIDE @cute.jit, so there was no MLIR context — causing 'Expected an MLIR object (got None)' in _pack_shape. Now ALL CuTe DSL calls happen INSIDE the @cute.jit function, matching fused_swiglu's pattern where __call__ is called from JIT context. Grid computation uses plain Python math (no MLIR needed). --- .../router/nvfp4_fused_router_kernel.py | 76 +++++++++---------- 1 file changed, 38 insertions(+), 38 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index d0b7d29b..308c1b11 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -282,49 +282,49 @@ class Nvfp4FusedRouterKernel: self.a_major_mode = a_major_mode self.b_major_mode = b_major_mode - tiled_mma = self._create_tiled_mma(a_dtype, a_major_mode, b_major_mode, sf_dtype) - tiled_mma_sfb = self._create_tiled_mma_sfb(a_dtype, a_major_mode, b_major_mode, sf_dtype) - self._setup_attributes(tiled_mma, tiled_mma_sfb, a_dtype, b_dtype, sf_dtype) - - # TMA atoms — following fused_swiglu.py exactly - # A - a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - a_op, mat_a, a_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - - # B - b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - b_op, mat_b, b_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - - # SFA - sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) - sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) - tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( - sfa_op, scale_a, sfa_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64) - - # SFB - sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) - sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( - sfb_op, scale_b, sfb_smem_layout, self.mma_tiler_sfb, tiled_mma_sfb, - self.cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Uint64) - # Grid: dense GEMM, one CTA per (M_tile, N_tile) - num_M_tiles = cute.ceil_div(M, self.cta_tile_shape_mnk[0]) - num_N_tiles = cute.ceil_div(N, self.cta_tile_shape_mnk[1]) - L = 1 + # Use Python math for grid calc (no MLIR needed) + cta_m = self.mma_tiler_mnk[0] # 128 + cta_n = self.mma_tiler_mnk[1] # 128 + num_M_tiles = (M + cta_m - 1) // cta_m + num_N_tiles = (N + cta_n - 1) // cta_n grid = (num_M_tiles * num_N_tiles, 1, 1) - tile_sched_params = utils.PersistentTileSchedulerParams( - (cutlass.Int32(num_M_tiles), cutlass.Int32(num_N_tiles), cutlass.Int32(L)), - (*self.cluster_shape_mn, 1)) - @cute.jit def _compiled_fn(mat_a, mat_b, scale_a, scale_b, e_bias, out_weights, out_ids): + # ALL CuTe DSL setup happens INSIDE JIT context + # This matches fused_swiglu's pattern where __call__ is called from JIT + tiled_mma = self._create_tiled_mma(a_dtype, a_major_mode, b_major_mode, sf_dtype) + tiled_mma_sfb = self._create_tiled_mma_sfb(a_dtype, a_major_mode, b_major_mode, sf_dtype) + self._setup_attributes(tiled_mma, tiled_mma_sfb, a_dtype, b_dtype, sf_dtype) + + # TMA atoms — following fused_swiglu.py exactly + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, mat_a, a_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) + + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, mat_b, b_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) + + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, scale_a, sfa_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64) + + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, scale_b, sfb_smem_layout, self.mma_tiler_sfb, tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Uint64) + + tile_sched_params = utils.PersistentTileSchedulerParams( + (cutlass.Int32(num_M_tiles), cutlass.Int32(num_N_tiles), cutlass.Int32(1)), + (cutlass.Int32(self.cluster_shape_mn[0]), cutlass.Int32(self.cluster_shape_mn[1]), cutlass.Int32(1))) + self._kernel( tiled_mma, tiled_mma_sfb, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b,