From 038fe81c689cbe34b7ea1a493302575b270518d1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 1 Jun 2026 15:03:54 +0000
Subject: [PATCH] Fix MoE non-fused L2 runtime gsa + update test harness for
 extra args

---
 dsv4/layers/moe.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py
index 4e937eb6..84c515ae 100644
--- a/dsv4/layers/moe.py
+++ b/dsv4/layers/moe.py
@@ -651,7 +651,11 @@ class Nvfp4MoE:
                 gate_silu = gate_silu.clamp(max=self._swiglu_limit)
                 up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
             activated = gate_silu * up
-        
+
+        # Compute runtime gsa for L2 from activated output (non-fused path)
+        if not self._fused_swiglu and getattr(self, '_use_runtime_gsa', False):
+            amax_l2 = activated.float().abs().max().clamp(min=1e-8).item()
+            self._l2_activation_global_scale = amax_l2 / (6.0 * 448.0)
         # === L2: down ===
         # Quantize activated (per-token) using GPU-only kernel, scatter into padded FP4 buffer.
         # For fused_swiglu path, slot_l2_x_fp4/sf already set by deinterleave_quantize_nvfp4_cuda.