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

Fix: properly remove sqrt(softplus) from CuTeDSL kernel

Browse Source 
Previous Python string replacement didn't match. Now using edit tool.
Kernel writes raw FP32 logits with gsa*gsb applied. sqrt(softplus)
is done in PyTorch after the kernel returns.
This commit is contained in:
biondizzle
2026-06-01 11:14:04 +00:00
parent 6efe90cd85
commit a4324781c3
1 changed files with 4 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
dsv4/kernels/router/nvfp4_fused_router_kernel.py
View File

@@ -712,21 +712,13 @@ class Nvfp4FusedRouterKernel:
acc_pipeline.consumer_release(acc_cs)
acc_cs.advance()
# Activation: sqrt(softplus(logit * gsa * gsb))
# Global scales are applied before the activation, same as
# how MoE epilogue applies them before SwiGLU.
# Apply global scale (gsa * gsb) to GEMM output
# The MMA output is (A * SFA) @ (B * SFB), missing gsa*gsb.
# Activation (sqrt(softplus)) is done in Python post-kernel
# because CuTeDSL MLIR crashes on exp+log+sqrt.
scale = cutlass.Float32(gsa * gsb)
acc_vec = tTR_rAcc.load()
for e in cutlass.range(cute.size(acc_vec), unroll=4):
logit = acc_vec[e] * scale
# softplus(x) = max(x, 0) + log(1 + exp(-|x|))
abs_x = cute.math.absf(logit)
pos = cute.math.fmax(logit, cutlass.Float32(0.0))
exp_neg = cute.math.exp(-abs_x)
sp = pos + cute.math.log(cutlass.Float32(1.0) + exp_neg)
acc_vec[e] = cute.math.sqrt(sp)
acc_vec = acc_vec * scale
tRS_rC.store(acc_vec.to(c_dtype))
# RMEM → SMEM