rocm: use numerically stable SiLU to avoid NaN on large-magnitude gate#498
Open
Flowers-of-Romance wants to merge 1 commit into
Open
rocm: use numerically stable SiLU to avoid NaN on large-magnitude gate#498Flowers-of-Romance wants to merge 1 commit into
Flowers-of-Romance wants to merge 1 commit into
Conversation
The ROCm MoE/FFN kernels computed SiLU as the naive `gate / (1.0f + expf(-gate))`. For a large negative gate, `expf(-gate)` overflows to +inf, so the expression becomes `-inf / inf = NaN`. The CPU reference (`ds4.c` `silu`/`sigmoid_stable`) deliberately avoids this by branching on the sign, so the two backends disagreed on extreme inputs and only ROCm produced NaN — which then propagates through the rest of the network (e.g. a later router selecting expert id -1). Add `dev_sigmoid_stable`/`dev_silu_stable` in `ds4_rocm_common.cuh` (matching the CPU sign-branched form) and replace the 27 naive SiLU sites across `ds4_rocm_moe.cuh`, `ds4_rocm_q8.cuh`, and `ds4_rocm_output.cuh`. No behavior change for well-conditioned inputs; only removes the overflow path on extreme gate values. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
OPS-NeoRetro
approved these changes
Jul 5, 2026
OPS-NeoRetro
left a comment
OPS-NeoRetro
left a comment
There was a problem hiding this comment.
So, your numerically stable Swish that Claude wrote is directly taken from ds4.c, but it's just ported to AMD ROCm. Are you planning to test it today or tomorrow? I hope that somebody would port this fix to CUDA, or I don't think that it's there.
Comment on lines
+11
to
+23
| /* Numerically stable sigmoid/SiLU matching the CPU reference (ds4.c | ||
| * sigmoid_stable/silu). The naive gate/(1+expf(-gate)) overflows for large | ||
| * negative gate: expf(-gate) -> +inf makes it -inf/inf = NaN, which the CPU | ||
| * path avoids. Use these everywhere the ROCm MoE/FFN applies SiLU. */ | ||
| __device__ __forceinline__ static float dev_sigmoid_stable(float x) { | ||
| if (x >= 0.0f) return 1.0f / (1.0f + expf(-x)); | ||
| const float e = expf(x); | ||
| return e / (1.0f + e); | ||
| } | ||
| __device__ __forceinline__ static float dev_silu_stable(float x) { | ||
| return x * dev_sigmoid_stable(x); | ||
| } | ||
|
|
There was a problem hiding this comment.
Great job. Claude did the common stable Swish.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The ROCm MoE/FFN kernels computed SiLU as the naive
gate / (1.0f + expf(-gate)). For a large negative gate,expf(-gate)overflows to +inf, so the expression becomes-inf / inf = NaN. The CPU reference (ds4.csilu/sigmoid_stable) deliberately avoids this by branching on the sign, so the two backends disagreed on extreme inputs and only ROCm produced NaN — which then propagates through the rest of the network (e.g. a later router selecting expert id -1).Add
dev_sigmoid_stable/dev_silu_stableinds4_rocm_common.cuh(matching the CPU sign-branched form) and replace the 27 naive SiLU sites acrossds4_rocm_moe.cuh,ds4_rocm_q8.cuh, andds4_rocm_output.cuh. No behavior change for well-conditioned inputs; only removes the overflow path on extreme gate values.