Add set_fused_swiglu() method to Nvfp4MoE

2026-06-02 07:59:57 +00:00
parent d8e17d70c1
commit 7904cf05c4
1 changed files with 4 additions and 0 deletions
--- 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).