From 6efe90cd85340868e29e3822f7d4892f3042b260 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 11:12:41 +0000 Subject: [PATCH] Move sqrt(softplus) out of CuTeDSL kernel into Python The CuTeDSL MLIR optimizer crashes (SIGABRT/core dump) on the combination of exp+log+sqrt in a for-range loop. The kernel now writes raw FP32 logits (with gsa*gsb applied) and sqrt(softplus) is done in PyTorch post-kernel. The GEMM is still pure NVFP4 Blackwell tensor cores. --- dsv4/kernels/router/nvfp4_fused_router_kernel.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index 1635c0b5..a44720f1 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -852,15 +852,20 @@ def run_nvfp4_fused_router( gsb=gsb_val, ) - # Add e_bias (selection bias) and run top-k - # The kernel writes sqrt(softplus(logits)) in FP32 - # activation_topk expects raw logits, so we pass the activated scores - # and tell it to skip the activation step + # Apply sqrt(softplus) activation in PyTorch (CuTeDSL MLIR crashes on exp+log+sqrt) + # softplus(x) = max(x, 0) + log(1 + exp(-|x|)) + abs_x = activated_scores.abs() + pos = activated_scores.clamp(min=0.0) + exp_neg = torch.exp(-abs_x) + sp = pos + torch.log1p(exp_neg) + activated = torch.sqrt(sp) + + # Top-k + renorm on activated scores from dsv4.kernels.router._activation_topk import run_fused_activation_topk_pre_activated out_weights = torch.empty(N, top_k, dtype=torch.float32, device=device) out_ids = torch.empty(N, top_k, dtype=torch.int32, device=device) run_fused_activation_topk_pre_activated( - activated_scores, e_bias, routed_scaling_factor, top_k, + activated, e_bias, routed_scaling_factor, top_k, out_weights, out_ids, )