biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: replace break statements with flag-based loops in router kernel (CuTeDSL restriction)

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:50:39 +00:00
parent 210391e571
commit d0d765e1f2
1 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
dsv4/kernels/router/dense_router_decode_kernel.py
View File

@@ -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