From d0d765e1f2ece94a13ae8be97fa0b55a3f827024 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 23:50:39 +0000 Subject: [PATCH] fix: replace break statements with flag-based loops in router kernel (CuTeDSL restriction) --- .../router/dense_router_decode_kernel.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/dsv4/kernels/router/dense_router_decode_kernel.py b/dsv4/kernels/router/dense_router_decode_kernel.py index 872a9f5c..2046875a 100644 --- a/dsv4/kernels/router/dense_router_decode_kernel.py +++ b/dsv4/kernels/router/dense_router_decode_kernel.py @@ -369,7 +369,8 @@ class DenseRouterDecodeKernel: # Sift down (k=6, fully unrolled) # Depth 0: children 1,2 root = 0 - while root < 3: + _done = cutlass.Bool(False) + while root < 3 and not _done: left = 2*root+1; right = 2*root+2 smallest = root if left < 6: @@ -379,8 +380,9 @@ class DenseRouterDecodeKernel: if hs[right] < hs[smallest] or (hs[right] == hs[smallest] and hi[right] > hi[smallest]): smallest = right if smallest == root: - break - ts = hs[root]; ti = hi[root]; ta = ha[root] + _done = cutlass.Bool(True) + if not _done: + ts = hs[root]; ti = hi[root]; ta = ha[root] hs[root] = hs[smallest]; hi[root] = hi[smallest]; ha[root] = ha[smallest] hs[smallest] = ts; hi[smallest] = ti; ha[smallest] = ta root = smallest @@ -410,7 +412,8 @@ class DenseRouterDecodeKernel: fs[0] = cs; fi[0] = ci; fa[0] = ca # Sift down r = 0 - while r < 3: + _done2 = cutlass.Bool(False) + while r < 3 and not _done2: l = 2*r+1; ri = 2*r+2; sm = r if l < 6: if fs[l] < fs[sm] or (fs[l] == fs[sm] and fi[l] > fi[sm]): @@ -418,8 +421,10 @@ class DenseRouterDecodeKernel: if ri < 6: if fs[ri] < fs[sm] or (fs[ri] == fs[sm] and fi[ri] > fi[sm]): sm = ri - if sm == r: break - ts=fs[r]; ti=fi[r]; ta=fa[r] + if sm == r: + _done2 = cutlass.Bool(True) + else: + ts=fs[r]; ti=fi[r]; ta=fa[r] fs[r]=fs[sm]; fi[r]=fi[sm]; fa[r]=fa[sm] fs[sm]=ts; fi[sm]=ti; fa[sm]=ta r = sm