From a4324781c3e8339cdd0afbb670b04b7c555208a5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 11:14:04 +0000 Subject: [PATCH] Fix: properly remove sqrt(softplus) from CuTeDSL kernel 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. --- dsv4/kernels/router/nvfp4_fused_router_kernel.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index a44720f1..df18139e 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -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