layernorm: FP32 SimplifiedLayerNorm/SkipSLN to prevent MoE routing collapse

src/fuse_attention.cpp

@@ -908,6 +908,116 @@ struct find_flash_decoding
     }
 };
 
+// Fuse the SkipSimplifiedLayerNorm (and SimplifiedLayerNorm) subgraph into a group op
+// before fuse_pointwise_reduce runs. This prevents fuse_pointwise from fusing the FP32
+// variance+rsqrt+multiply chain into a kernel that reverts the gamma multiply to FP16.
+//
+// The SLN subgraph emitted by parse_skip_simplified_layer_normalization.cpp is:
+//   add(x, skip) → convert(→fp32) → mul → mul → reduce_mean(rms, fp32)
+//   → add(rms, eps_fp32) → rsqrt(fp32) → mul(float_x, rrms, fp32)
+//   → mul(result, gamma_fp32, fp32) → convert(→fp16) [= result output]
+//
+// The matcher identifies the result convert(→fp16) whose input traces through rsqrt(fp32).
+// This is sufficient to uniquely identify SLN subgraphs (rsqrt in FP32 only appears here).
+struct find_skip_simplified_layer_norm
+{
+    std::size_t* counter;
+
+    auto matcher() const
+    {
+        // Match: convert(→fp16) whose input is mul whose input (transitively) is rsqrt(fp32)
+        auto rsqrt_fp32 = match::name("rsqrt")(
+            match::arg(0)(match::any().bind("rms_ep")));           // rms + eps (fp32)
+        auto mul_x_rrms = match::name("mul")(                      // float_x * rrms_fp32
+            match::any_arg(0, 1)(rsqrt_fp32));
+        auto mul_gamma  = match::name("mul")(                      // result_fp32 * gamma_fp32
+            match::any_arg(0, 1)(mul_x_rrms));
+        // Final convert back to fp16 / io_dtype
+        return match::name("convert")(
+            match::arg(0)(mul_gamma)).bind("sln_result");
+    }
+
+    std::string get_count() const { return std::to_string((*counter)++); }
+
+    void apply(module_pass_manager& mpm, const match::matcher_result& r) const
+    {
+        auto sln_result = r.instructions["sln_result"];
+        auto rms_ep     = r.instructions["rms_ep"];
+
+        // Collect the SLN subgraph: all instructions between rms_ep (inclusive) and
+        // sln_result (inclusive) that are exclusively part of this SLN chain.
+        // Walk from rms_ep forward to sln_result, collecting ops.
+        auto& m = mpm.get_module();
+
+        // Gather all instructions reachable from rms_ep that lead to sln_result
+        std::unordered_set<instruction_ref> sln_inss;
+        auto is_sln_op = [](instruction_ref ins) {
+            if(ins->name() == "contiguous")
+                return true;
+            if(ins->get_operator().attributes().get("pointwise", false))
+                return true;
+            static const std::unordered_set<std::string> sln_ops = {
+                "mul", "add", "sub", "div", "convert", "reduce_mean", "reduce_sum",
+                "multibroadcast", "broadcast", "rsqrt", "reshape", "squeeze", "unsqueeze"};
+            return contains(sln_ops, ins->name());
+        };
+
+        std::function<bool(instruction_ref)> collect = [&](instruction_ref ins) -> bool {
+            if(contains(sln_inss, ins))
+                return true;
+            if(ins == rms_ep)
+            {
+                sln_inss.insert(ins);
+                return true;
+            }
+            if(not is_sln_op(ins))
+                return false;
+            bool any_input_in = false;
+            for(auto input : ins->inputs())
+            {
+                if(collect(input))
+                    any_input_in = true;
+            }
+            if(any_input_in)
+                sln_inss.insert(ins);
+            return any_input_in;
+        };
+        collect(sln_result);
+
+        if(sln_inss.size() < 3) // sanity: need at least rsqrt + mul + convert
+            return;
+
+        // Sort topologically
+        std::vector<instruction_ref> sorted_inss(sln_inss.begin(), sln_inss.end());
+        std::sort(sorted_inss.begin(), sorted_inss.end(), [&](instruction_ref x, instruction_ref y) {
+            return std::distance(m.begin(), x) < std::distance(m.begin(), y);
+        });
+
+        // Build submodule
+        module m_sln;
+        std::unordered_map<instruction_ref, instruction_ref> map_mm_to_msln;
+        m_sln.fuse(sorted_inss, &map_mm_to_msln);
+        dead_code_elimination{}.apply(m_sln);
+
+        m_sln.add_return({map_mm_to_msln.at(sln_result)});
+
+        auto map_msln_to_mm = [&] {
+            std::unordered_map<instruction_ref, instruction_ref> inv;
+            for(auto& [k, v] : map_mm_to_msln)
+                inv[v] = k;
+            return inv;
+        }();
+        auto new_inputs = m_sln.get_inputs(map_msln_to_mm);
+
+        module_ref mpm_sln = mpm.create_module("sln" + get_count(), std::move(m_sln));
+        mpm_sln->set_bypass();
+
+        auto group_ins = m.insert_instruction(
+            sln_result, make_op("group", {{"tag", "skip_layer_norm"}}), new_inputs, {mpm_sln});
+        m.replace_instruction(sln_result, group_ins);
+    }
+};
+
 struct find_kv_cache_attention
 {
     std::size_t* counter;
@@ -1109,6 +1219,13 @@ void fuse_attention::apply(module_pass_manager& mpm) const
 {
     std::size_t counter = 0;
 
+    // Fuse SkipSimplifiedLayerNorm into a group op (opaque to fuse_pointwise) so that
+    // the FP32 variance+rsqrt+mul chain is not broken by fuse_pointwise coercing the
+    // gamma multiply back to FP16. Must run before fuse_pointwise_reduce.
+    match::find_matches(mpm, find_skip_simplified_layer_norm{.counter = &counter});
+    mpm.get_module().sort();
+    mpm.run_pass(dead_code_elimination{});
+
     // Fuse kv-cache attention by default
     match::find_matches(mpm, find_kv_cache_attention{.counter = &counter});
     mpm.get_module().sort();

src/onnx/parse_simplified_layer_normalization.cpp

@@ -86,15 +86,21 @@ struct parse_simplified_layer_normalization : op_parser<parse_simplified_layer_n
             make_op("convert", {{"target_type", migraphx::shape::float_type}}), x);
         auto x_sq = info.add_common_op("mul", float_x, float_x);
         auto rms  = info.add_instruction(make_op("reduce_mean", {{"axes", {axis}}}), x_sq);
-        rms       = info.add_instruction(make_op("convert", {{"target_type", x_dtype}}), rms);
-        auto mean = rms;
+        // Keep variance in FP32 through rsqrt (same fix as SkipSimplifiedLayerNorm).
+        auto mean = info.add_instruction(
+            make_op("convert", {{"target_type", x_dtype}}), rms);
         epsilon =
             (x_dtype == migraphx::shape::half_type and std::abs(epsilon) < 1e-7) ? 1e-7 : epsilon;
-        auto eps    = info.add_literal(migraphx::literal{migraphx::shape{x_dtype}, {epsilon}});
-        rms         = info.add_common_op("add", rms, eps);
-        auto rrms   = info.add_instruction(make_op("rsqrt"), rms);
-        auto result = info.add_common_op("mul", x, rrms);
-        result      = info.add_common_op("mul", result, scale);
+        auto eps_f32  = info.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::float_type}, {epsilon}});
+        auto rms_ep   = info.add_common_op("add", rms, eps_f32);
+        auto rrms_f32 = info.add_instruction(make_op("rsqrt"), rms_ep);        // FP32
+        auto scale_f32 = info.add_instruction(
+            make_op("convert", {{"target_type", migraphx::shape::float_type}}), scale);
+        scale_f32 = info.add_instruction(make_op("contiguous"), scale_f32);
+        auto result_f32 = info.add_common_op("mul", float_x, rrms_f32);
+        result_f32      = info.add_common_op("mul", result_f32, scale_f32);
+        auto rrms   = info.add_instruction(make_op("convert", {{"target_type", x_dtype}}), rrms_f32);
+        auto result = info.add_instruction(make_op("convert", {{"target_type", x_dtype}}), result_f32);
 
         return {result, mean, rrms};
     }

src/onnx/parse_skip_simplified_layer_normalization.cpp

@@ -108,15 +108,32 @@ struct parse_skip_simplified_layer_normalization
             make_op("convert", {{"target_type", migraphx::shape::float_type}}), x);
         auto x_sq = info.add_common_op("mul", float_x, float_x);
         auto rms  = info.add_instruction(make_op("reduce_mean", {{"axes", {axis}}}), x_sq);
-        rms       = info.add_instruction(make_op("convert", {{"target_type", x_dtype}}), rms);
-        auto mean = rms;
+        // Full FP32 normalization: mirror the DML reference (ComputeSkipSLNCPU in
+        // dml/hip_qmoe/qmoe_hip_combined_op.cpp): mean_sq in FP32, inv_std in FP32,
+        // x*inv_std*gamma all in FP32, only the final output cast back to io_dtype.
+        // Converting rrms or intermediate results to FP16 early reintroduces the
+        // precision loss that causes router logits to drift 3-10x by layer 20.
+        auto mean = info.add_instruction(
+            make_op("convert", {{"target_type", x_dtype}}), rms); // FP16 mean for output only
         epsilon =
             (x_dtype == migraphx::shape::half_type and std::abs(epsilon) < 1e-7) ? 1e-7 : epsilon;
-        auto eps    = info.add_literal(migraphx::literal{migraphx::shape{x_dtype}, {epsilon}});
-        rms         = info.add_common_op("add", rms, eps);
-        auto rrms   = info.add_instruction(make_op("rsqrt"), rms);
-        auto result = info.add_common_op("mul", x, rrms);
-        result      = info.add_common_op("mul", result, gamma);
+        auto eps_f32  = info.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::float_type}, {epsilon}});
+        auto rms_ep   = info.add_common_op("add", rms, eps_f32);         // FP32
+        auto rrms_f32 = info.add_instruction(make_op("rsqrt"), rms_ep);  // FP32
+        // Cast gamma to FP32 so mul stays in FP32
+        // Use contiguous to anchor the FP32 value and prevent eliminate_convert from
+        // removing the cast when gamma is used elsewhere in FP16.
+        auto gamma_f32 = info.add_instruction(
+            make_op("convert", {{"target_type", migraphx::shape::float_type}}), gamma);
+        gamma_f32 = info.add_instruction(make_op("contiguous"), gamma_f32);
+        // Compute x*rrms*gamma entirely in FP32 (add_common_op handles broadcasting)
+        auto result_f32 = info.add_common_op("mul", float_x, rrms_f32);
+        result_f32      = info.add_common_op("mul", result_f32, gamma_f32);
+        // Cast final result back to io_dtype
+        auto rrms = info.add_instruction(
+            make_op("convert", {{"target_type", x_dtype}}), rrms_f32); // kept for output slot
+        auto result = info.add_instruction(
+            make_op("convert", {{"target_type", x_dtype}}), result_f32);
         if(args.size() == 4)
         {
             result = info.add_common_op("add", result, bias);

src/targets/gpu/fuse_mlir.cpp

@@ -1618,6 +1618,27 @@ void fuse_mlir::apply(module_pass_manager& mpm) const
     match::find_matches(mpm, find_mlir_kv_cache_attention_op{mlir_mode::all});
     mpm.run_pass(dead_code_elimination{});
 
+    // Compile SkipSimplifiedLayerNorm group (tagged 'skip_layer_norm') via MLIR.
+    // This group was created by find_skip_simplified_layer_norm in fuse_attention
+    // to prevent fuse_pointwise from reverting the FP32 gamma multiply to FP16.
+    // MLIR compiles the submodule as a single kernel respecting the op types exactly.
+    struct find_mlir_skip_sln_op
+    {
+        auto matcher() const
+        {
+            return match::name("group")(match::has_op_value("tag", "skip_layer_norm"));
+        }
+        void apply(module_pass_manager& mpm, const match::matcher_result& r) const
+        {
+            auto group   = r.result;
+            auto* m_sln  = group->module_inputs()[0];
+            mpm.get_module().replace_instruction(
+                group, mlir_op{group->get_operator()}, mlir_contiguous(mpm, group->inputs()), {m_sln});
+        }
+    };
+    match::find_matches(mpm, find_mlir_skip_sln_op{});
+    mpm.run_pass(dead_code_elimination{});
+
     match::find_matches(mpm, find_mlir_attention_op{});
     mpm.run_pass(dead_code_elimination{});
 

src/targets/gpu/kernels/include/migraphx/kernels/skip_simplified_layer_norm.hpp

@@ -0,0 +1,109 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+/*
+ * SkipSimplifiedLayerNorm FP32 kernel for GPT-OSS-20B correctness.
+ *
+ * Mirrors the DML validated reference (ComputeSkipSLNCPU in dml/hip_qmoe/
+ * qmoe_hip_combined_op.cpp, lines 2178-2263): mean_sq in FP32, inv_std in FP32,
+ * x*inv_std*gamma all in FP32. Only the output is converted to fp16.
+ *
+ * Root cause this fixes: MIGraphX's fuse_pointwise_reduce merges the SLN ops
+ * into a single kernel where the gamma multiply reverts to fp16 because gamma
+ * is a shared fp16 weight tensor. This custom kernel keeps everything in FP32.
+ *
+ * Kernel: one block per token, blockDim.x threads reduce hidden_size elements.
+ * Uses warp shuffle for the variance reduction.
+ */
+#ifndef MIGRAPHX_GUARD_KERNELS_SKIP_SIMPLIFIED_LAYER_NORM_HPP
+#define MIGRAPHX_GUARD_KERNELS_SKIP_SIMPLIFIED_LAYER_NORM_HPP
+
+#include <migraphx/kernels/index.hpp>
+#include <migraphx/kernels/tensor_view.hpp>
+
+namespace migraphx {
+
+/*
+ * skip_simplified_layer_norm<BLOCK_SIZE>(input, skip, gamma, output, eps, hidden_size)
+ *
+ * All of mean_sq, inv_std, x*inv_std*gamma computed in FP32.
+ * Input/output in fp16; gamma in fp16 (converted to fp32 inside kernel).
+ *
+ * Launch: gridDim.x = num_tokens, blockDim.x = BLOCK_SIZE (e.g. 256)
+ */
+template <index_int BLOCK_SIZE, class Input, class Skip, class Gamma, class Output>
+__device__ void skip_simplified_layer_norm(const Input  input,
+                                           const Skip   skip,
+                                           const Gamma  gamma,
+                                           Output       output,
+                                           float        eps,
+                                           index_int    hidden_size)
+{
+    const index_int token_idx  = blockIdx.x;
+    const index_int thread_idx = threadIdx.x;
+
+    // Shared memory for warp-level reduction of sum_sq
+    __shared__ float shmem[BLOCK_SIZE / 32]; // one slot per warp
+
+    // Each thread accumulates sum of (x+skip)^2 over its chunk
+    float local_sum_sq = 0.0f;
+    for(index_int i = thread_idx; i < hidden_size; i += BLOCK_SIZE)
+    {
+        float x_val  = __half2float(input[token_idx * hidden_size + i]);
+        float sk_val = __half2float(skip[token_idx * hidden_size + i]);
+        float v      = x_val + sk_val;
+        local_sum_sq += v * v;
+    }
+
+    // Warp-level reduction
+    for(int offset = 16; offset > 0; offset >>= 1)
+        local_sum_sq += __shfl_xor(local_sum_sq, offset);
+
+    if((thread_idx & 31) == 0) // lane 0 of each warp writes to shared
+        shmem[thread_idx >> 5] = local_sum_sq;
+    __syncthreads();
+
+    // Block-level reduction in first warp
+    float block_sum_sq = 0.0f;
+    if(thread_idx < (BLOCK_SIZE / 32))
+        block_sum_sq = shmem[thread_idx];
+    for(int offset = (BLOCK_SIZE / 64); offset > 0; offset >>= 1)
+        block_sum_sq += __shfl_xor(block_sum_sq, offset);
+
+    // Broadcast inv_std to all threads
+    float inv_std = __shfl(1.0f / sqrtf(block_sum_sq / (float)hidden_size + eps), 0);
+
+    // Each thread applies normalization in FP32 and writes fp16 output
+    for(index_int i = thread_idx; i < hidden_size; i += BLOCK_SIZE)
+    {
+        float x_val  = __half2float(input[token_idx * hidden_size + i]);
+        float sk_val = __half2float(skip[token_idx * hidden_size + i]);
+        float g_val  = __half2float(gamma[i]); // gamma converted to FP32 here
+        float result = (x_val + sk_val) * inv_std * g_val;
+        output[token_idx * hidden_size + i] = __float2half(result);
+    }
+}
+
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_KERNELS_SKIP_SIMPLIFIED_LAYER_NORM_HPP