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.
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user