diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py index 82b6ac58..6d7bf149 100644 --- a/dsv4/layers/moe.py +++ b/dsv4/layers/moe.py @@ -104,6 +104,10 @@ class Nvfp4MoE: """Set the swiglu_limit for activation clamping.""" self._swiglu_limit = limit + def set_fused_swiglu(self, enabled: bool): + """Enable fused L1 GEMM + SwiGLU kernel (saves 240+ BF16 kernel launches per token).""" + self._fused_swiglu = enabled + def _fill_token_indices(self): """Fill _token_indices with [0,0,..0, 1,1,..1, ...] (each token repeated top_k times).