From f33ca41c2a1a8772be8a5240b12107f6ea086c7a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 09:13:53 +0000 Subject: [PATCH] Fused router: replace nested if/else top-k with flat find-min-replace approach The 5-level nested if/else for sorted insertion created O(2^5) MLIR regions that crashed the CuTeDSL MLIR optimizer (SIGABRT). New approach: - Find-min-replace: scan 6 entries to find minimum (sequential, 1-level nesting) - Replace the minimum if new score > min (flat conditionals by index) - Selection sort the final 6 entries after SMEM merge (descending order) - All conditionals are FLAT (at most 1 level of nesting) This should avoid the MLIR optimizer explosion while producing identical results. --- .../router/nvfp4_fused_router_kernel.py | 238 +++++++++++++----- 1 file changed, 180 insertions(+), 58 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index 097c9186..6ab2871a 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -685,15 +685,16 @@ class Nvfp4FusedRouterKernel: # EPILOGUE WARPS — TMEM → registers → router logic → GMEM # ============================================================ # - # Strategy: + # Strategy (FLAT top-k — no nested if/else): # 1. Read TMEM accumulator into registers via paired t2r copy - # (using epilogue_tmem_copy_and_partition from CUTLASS) # 2. For each element: compute act = sqrt(softplus(logit)), # score = act + e_bias[expert_idx] - # 3. Maintain per-thread running top-k via sorted insertion - # (fully unrolled for top_k=6, descending order) + # 3. Maintain per-thread running top-k via find-min-replace: + # - Find the minimum score among the current top-k (flat sequential scan) + # - If new score > min, replace the min entry (flat conditional by index) + # This avoids the 5-level nested if/else that crashes the MLIR optimizer. # 4. After all tiles: write local top-k to SMEM, - # thread 0 merges, sorts, renormalizes, writes to GMEM + # thread 0 merges using same flat approach, renormalizes, writes to GMEM # if warp_idx in self.epilogue_warp_id: if cute.size(self.cluster_shape_mn) > 1: @@ -710,12 +711,15 @@ class Nvfp4FusedRouterKernel: tCtAcc_base, epi_tile, self.epilogue_warp_id, acc_dtype, use_2cta) tTR_rAcc = tiled_copy_t2r.fragments_slice(tiled_copy_t2r, tTR_tAcc_base) - # Per-thread running top-k — individual scalar variables - # Stored in DESCENDING order: s0 >= s1 >= s2 >= s3 >= s4 >= s5 + # Per-thread running top-k (UNsorted — find-min-replace strategy) + # These hold the 6 best (score, expert_index, unbiased_activation) seen so far. + # They are NOT maintained in sorted order during accumulation. + # Sorting is done once at the end by thread 0 after the SMEM merge. TK = self.top_k - s0 = cutlass.Float32(-1e30); s1 = cutlass.Float32(-1e30) - s2 = cutlass.Float32(-1e30); s3 = cutlass.Float32(-1e30) - s4 = cutlass.Float32(-1e30); s5 = cutlass.Float32(-1e30) + NEG_INF = cutlass.Float32(-1e30) + s0 = NEG_INF; s1 = NEG_INF + s2 = NEG_INF; s3 = NEG_INF + s4 = NEG_INF; s5 = NEG_INF i0 = cutlass.Int32(-1); i1 = cutlass.Int32(-1) i2 = cutlass.Int32(-1); i3 = cutlass.Int32(-1) i4 = cutlass.Int32(-1); i5 = cutlass.Int32(-1) @@ -741,7 +745,6 @@ class Nvfp4FusedRouterKernel: # Get tile N offset (which expert slice this tile covers) tc = wt.tile_idx tile_n_offset = tc[1] * self.cta_tile_shape_mnk[1] - tile_m_base = tc[0] // cute.size(tiled_mma.thr_id.shape) * self.cta_tile_shape_mnk[0] tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) @@ -782,29 +785,40 @@ class Nvfp4FusedRouterKernel: # score = act + e_bias (for selection only) score = act + e_bias_tensor[e_idx] - # Sorted insertion into descending top-6 - # s0 >= s1 >= s2 >= s3 >= s4 >= s5 - if score > s5: - if score > s4: - s5 = s4; i5 = i4; a5 = a4 - if score > s3: - s4 = s3; i4 = i3; a4 = a3 - if score > s2: - s3 = s2; i3 = i2; a3 = a2 - if score > s1: - s2 = s1; i2 = i1; a2 = a1 - if score > s0: - s1 = s0; i1 = i0; a1 = a0 - s0 = score; i0 = e_idx; a0 = act - else: - s1 = score; i1 = e_idx; a1 = act - else: - s2 = score; i2 = e_idx; a2 = act - else: - s3 = score; i3 = e_idx; a3 = act - else: - s4 = score; i4 = e_idx; a4 = act - else: + # FLAT top-k: find minimum, then replace if new > min + # Step 1: find the minimum score among s0..s5 + min_s = s0 + min_k = cutlass.Int32(0) + if s1 < min_s: + min_s = s1 + min_k = cutlass.Int32(1) + if s2 < min_s: + min_s = s2 + min_k = cutlass.Int32(2) + if s3 < min_s: + min_s = s3 + min_k = cutlass.Int32(3) + if s4 < min_s: + min_s = s4 + min_k = cutlass.Int32(4) + if s5 < min_s: + min_s = s5 + min_k = cutlass.Int32(5) + + # Step 2: if new score > minimum, replace the min entry + if score > min_s: + # Replace at position min_k (flat conditionals, NOT nested) + if min_k == cutlass.Int32(0): + s0 = score; i0 = e_idx; a0 = act + if min_k == cutlass.Int32(1): + s1 = score; i1 = e_idx; a1 = act + if min_k == cutlass.Int32(2): + s2 = score; i2 = e_idx; a2 = act + if min_k == cutlass.Int32(3): + s3 = score; i3 = e_idx; a3 = act + if min_k == cutlass.Int32(4): + s4 = score; i4 = e_idx; a4 = act + if min_k == cutlass.Int32(5): s5 = score; i5 = e_idx; a5 = act # Release accumulator (non-overlapping case) @@ -831,14 +845,14 @@ class Nvfp4FusedRouterKernel: epi_bar.arrive_and_wait() - # Thread 0 of warp 0 does the final merge + store + # Thread 0 of warp 0 does the final merge + sort + store if warp_idx == 0 and tidx == 0: # Initialize final top-6 from thread 0's data fs0 = s0; fs1 = s1; fs2 = s2; fs3 = s3; fs4 = s4; fs5 = s5 fi0 = i0; fi1 = i1; fi2 = i2; fi3 = i3; fi4 = i4; fi5 = i5 fa0 = a0; fa1 = a1; fa2 = a2; fa3 = a3; fa4 = a4; fa5 = a5 - # Merge all other threads (1..127) + # Merge all other threads (1..127) using flat find-min-replace for t in cutlass.range(1, 128, unroll=1): for k_idx in cutlass.range(TK, unroll=1): cs = storage.merge_scores.data_ptr()[t * TK + k_idx] @@ -846,30 +860,138 @@ class Nvfp4FusedRouterKernel: ca = storage.merge_acts.data_ptr()[t * TK + k_idx] # Only merge if this is a valid entry (index >= 0) if ci >= cutlass.Int32(0): - # Sorted insertion into final top-6 (descending) - if cs > fs5: - if cs > fs4: - fs5 = fs4; fi5 = fi4; fa5 = fa4 - if cs > fs3: - fs4 = fs3; fi4 = fi3; fa4 = fa3 - if cs > fs2: - fs3 = fs2; fi3 = fi2; fa3 = fa2 - if cs > fs1: - fs2 = fs1; fi2 = fi1; fa2 = fa1 - if cs > fs0: - fs1 = fs0; fi1 = fi0; fa1 = fa0 - fs0 = cs; fi0 = ci; fa0 = ca - else: - fs1 = cs; fi1 = ci; fa1 = ca - else: - fs2 = cs; fi2 = ci; fa2 = ca - else: - fs3 = cs; fi3 = ci; fa3 = ca - else: - fs4 = cs; fi4 = ci; fa4 = ca - else: + # Find minimum among final top-6 + fmin_s = fs0 + fmin_k = cutlass.Int32(0) + if fs1 < fmin_s: + fmin_s = fs1 + fmin_k = cutlass.Int32(1) + if fs2 < fmin_s: + fmin_s = fs2 + fmin_k = cutlass.Int32(2) + if fs3 < fmin_s: + fmin_s = fs3 + fmin_k = cutlass.Int32(3) + if fs4 < fmin_s: + fmin_s = fs4 + fmin_k = cutlass.Int32(4) + if fs5 < fmin_s: + fmin_s = fs5 + fmin_k = cutlass.Int32(5) + # Replace if candidate is better + if cs > fmin_s: + if fmin_k == cutlass.Int32(0): + fs0 = cs; fi0 = ci; fa0 = ca + if fmin_k == cutlass.Int32(1): + fs1 = cs; fi1 = ci; fa1 = ca + if fmin_k == cutlass.Int32(2): + fs2 = cs; fi2 = ci; fa2 = ca + if fmin_k == cutlass.Int32(3): + fs3 = cs; fi3 = ci; fa3 = ca + if fmin_k == cutlass.Int32(4): + fs4 = cs; fi4 = ci; fa4 = ca + if fmin_k == cutlass.Int32(5): fs5 = cs; fi5 = ci; fa5 = ca + # Selection sort the final 6 entries into descending order + # 6 passes: each finds the maximum among remaining entries + # Pass 0: find max among all 6 → position 0 + # (Using pairwise fmax to avoid nested if/else) + # + # For top_k=6, selection sort with flat max-finding: + # max(a,b) via cute.math.fmax, then compare to find index + # + # Since the top-k is only 6 entries, we can do this + # with a simple approach: for each position, scan all + # remaining entries to find the best. + # + # We'll sort IN-PLACE by repeatedly finding the max of + # the remaining tail and swapping it to the current position. + + # Pass 0: max of [0..5] + m0_s = fs0; m0_i = fi0; m0_a = fa0; m0_k = cutlass.Int32(0) + if fs1 > m0_s: + m0_s = fs1; m0_i = fi1; m0_a = fa1; m0_k = cutlass.Int32(1) + if fs2 > m0_s: + m0_s = fs2; m0_i = fi2; m0_a = fa2; m0_k = cutlass.Int32(2) + if fs3 > m0_s: + m0_s = fs3; m0_i = fi3; m0_a = fa3; m0_k = cutlass.Int32(3) + if fs4 > m0_s: + m0_s = fs4; m0_i = fi4; m0_a = fa4; m0_k = cutlass.Int32(4) + if fs5 > m0_s: + m0_s = fs5; m0_i = fi5; m0_a = fa5; m0_k = cutlass.Int32(5) + # Swap position 0 with the max (flat conditionals by position) + if m0_k == cutlass.Int32(1): + fs1 = t_s; fi1 = t_i; fa1 = t_a + if m0_k == cutlass.Int32(2): + fs2 = t_s; fi2 = t_i; fa2 = t_a + if m0_k == cutlass.Int32(3): + fs3 = t_s; fi3 = t_i; fa3 = t_a + if m0_k == cutlass.Int32(4): + fs4 = t_s; fi4 = t_i; fa4 = t_a + if m0_k == cutlass.Int32(5): + fs5 = t_s; fi5 = t_i; fa5 = t_a + # (if m0_k == 0, swap is a no-op) + + # Pass 1: max of [1..5] + m1_s = fs1; m1_i = fi1; m1_a = fa1; m1_k = cutlass.Int32(1) + if fs2 > m1_s: + m1_s = fs2; m1_i = fi2; m1_a = fa2; m1_k = cutlass.Int32(2) + if fs3 > m1_s: + m1_s = fs3; m1_i = fi3; m1_a = fa3; m1_k = cutlass.Int32(3) + if fs4 > m1_s: + m1_s = fs4; m1_i = fi4; m1_a = fa4; m1_k = cutlass.Int32(4) + if fs5 > m1_s: + m1_s = fs5; m1_i = fi5; m1_a = fa5; m1_k = cutlass.Int32(5) + t_s = fs1; t_i = fi1; t_a = fa1 + fs1 = m1_s; fi1 = m1_i; fa1 = m1_a + if m1_k == cutlass.Int32(2): + fs2 = t_s; fi2 = t_i; fa2 = t_a + if m1_k == cutlass.Int32(3): + fs3 = t_s; fi3 = t_i; fa3 = t_a + if m1_k == cutlass.Int32(4): + fs4 = t_s; fi4 = t_i; fa4 = t_a + if m1_k == cutlass.Int32(5): + fs5 = t_s; fi5 = t_i; fa5 = t_a + + # Pass 2: max of [2..5] + m2_s = fs2; m2_i = fi2; m2_a = fa2; m2_k = cutlass.Int32(2) + if fs3 > m2_s: + m2_s = fs3; m2_i = fi3; m2_a = fa3; m2_k = cutlass.Int32(3) + if fs4 > m2_s: + m2_s = fs4; m2_i = fi4; m2_a = fa4; m2_k = cutlass.Int32(4) + if fs5 > m2_s: + m2_s = fs5; m2_i = fi5; m2_a = fa5; m2_k = cutlass.Int32(5) + t_s = fs2; t_i = fi2; t_a = fa2 + fs2 = m2_s; fi2 = m2_i; fa2 = m2_a + if m2_k == cutlass.Int32(3): + fs3 = t_s; fi3 = t_i; fa3 = t_a + if m2_k == cutlass.Int32(4): + fs4 = t_s; fi4 = t_i; fa4 = t_a + if m2_k == cutlass.Int32(5): + fs5 = t_s; fi5 = t_i; fa5 = t_a + + # Pass 3: max of [3..5] + m3_s = fs3; m3_i = fi3; m3_a = fa3; m3_k = cutlass.Int32(3) + if fs4 > m3_s: + m3_s = fs4; m3_i = fi4; m3_a = fa4; m3_k = cutlass.Int32(4) + if fs5 > m3_s: + m3_s = fs5; m3_i = fi5; m3_a = fa5; m3_k = cutlass.Int32(5) + t_s = fs3; t_i = fi3; t_a = fa3 + fs3 = m3_s; fi3 = m3_i; fa3 = m3_a + if m3_k == cutlass.Int32(4): + fs4 = t_s; fi4 = t_i; fa4 = t_a + if m3_k == cutlass.Int32(5): + fs5 = t_s; fi5 = t_i; fa5 = t_a + + # Pass 4: max of [4..5] + if fs5 > fs4: + t_s = fs4; t_i = fi4; t_a = fa4 + fs4 = fs5; fi4 = fi5; fa4 = fa5 + fs5 = t_s; fi5 = t_i; fa5 = t_a + # Pass 5: [5] is alone — nothing to do + + # Now fs0..fs5 are in descending order. # Renormalize: w = act / sum(act) * scaling act_sum = fa0 + fa1 + fa2 + fa3 + fa4 + fa5 inv_sum = cutlass.Float32(1.0) / act_sum