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, )