Add BF16 fallback for shared expert: dequantize NVFP4 → BF16 F.linear

Same approach as router gate. enable_bf16_fallback() pre-materializes
BF16 weight matrices from NVFP4 weights, then _run_impl uses F.linear
instead of NVFP4 GEMM.

dsv4/layers/shared_expert.py

@@ -101,6 +101,10 @@ class Nvfp4SharedExpert:
         self._expert_offsets_buf = None
         self._buffers_allocated = False
 
+        self._use_bf16_fallback = False  # Set True to use BF16 F.linear instead of NVFP4 GEMM
+        self._bf16_l1_weight = None
+        self._bf16_l2_weight = None
+
     def set_swiglu_limit(self, limit: float):
         self.swiglu_limit = limit
 
@@ -108,6 +112,31 @@ class Nvfp4SharedExpert:
         """Enable fused L1 GEMM + SwiGLU kernel (1-group variant of MoE fused kernel)."""
         self._fused_swiglu = enabled
 
+    def enable_bf16_fallback(self):
+        """Dequantize NVFP4 weights to BF16 and use F.linear instead of NVFP4 GEMM."""
+        self._use_bf16_fallback = True
+        self._fused_swiglu = False  # fused kernel not compatible with BF16 path
+        from dsv4.ops.quantize import dequantize_nvfp4 as _dequant
+        # L1: (2*intermediate, hidden) — gate+up interleaved
+        if self.l1_fp4 is not None and len(self.l1_fp4) > 0 and self.l1_fp4[0] is not None:
+            l1_w = self.l1_fp4[0]
+            l1_s = self.l1_sf[0] if self.l1_sf is not None and len(self.l1_sf) > 0 else None
+            l1_gsb = self.l1_gs[0] if self.l1_gs is not None and len(self.l1_gs) > 0 else 1.0
+            # Fold weight_scale_2 into gsb
+            if self.l1_ws2 is not None and len(self.l1_ws2) > 0 and self.l1_ws2[0] is not None:
+                l1_gsb = l1_gsb * self.l1_ws2[0].float().item()
+            gsa_l1 = torch.tensor([l1_gsb] * l1_w.shape[0], device=l1_w.device, dtype=torch.float32)
+            self._bf16_l1_weight = _dequant(l1_w, l1_s, gsa_l1)
+        # L2: (hidden, intermediate) — down projection
+        if self.l2_fp4 is not None and len(self.l2_fp4) > 0 and self.l2_fp4[0] is not None:
+            l2_w = self.l2_fp4[0]
+            l2_s = self.l2_sf[0] if self.l2_sf is not None and len(self.l2_sf) > 0 else None
+            l2_gsb = self.l2_gs[0] if self.l2_gs is not None and len(self.l2_gs) > 0 else 1.0
+            if self.l2_ws2 is not None and len(self.l2_ws2) > 0 and self.l2_ws2[0] is not None:
+                l2_gsb = l2_gsb * self.l2_ws2[0].float().item()
+            gsa_l2 = torch.tensor([l2_gsb] * l2_w.shape[0], device=l2_w.device, dtype=torch.float32)
+            self._bf16_l2_weight = _dequant(l2_w, l2_s, gsa_l2)
+
     def finalize_weights(self):
         """Process weights for CuTeDSL GEMM. Must be called after setting l1/l2 weights."""
         # Convert uint8 checkpoint weights to float4_e2m1fn_x2 view
@@ -388,6 +417,18 @@ class Nvfp4SharedExpert:
 
     def _run_impl(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Actual implementation — called via custom autograd to be torch.compile-safe."""
+        if self._use_bf16_fallback and self._bf16_l1_weight is not None:
+            # BF16 path: F.linear with dequantized weights
+            l1_out = torch.nn.functional.linear(hidden_states, self._bf16_l1_weight)
+            gate = l1_out[:, :self.intermediate_size]
+            up = l1_out[:, self.intermediate_size:]
+            if self.swiglu_limit is not None:
+                gate = gate.clamp(max=self.swiglu_limit)
+                up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit)
+            intermediate = torch.nn.functional.silu(gate) * up
+            output = torch.nn.functional.linear(intermediate, self._bf16_l2_weight)
+            return output
+
         self._ensure_initialized()
 
         if self._fused_swiglu:

single_shot_inference.py

@@ -1359,18 +1359,9 @@ def main():
         se.set_fused_swiglu(True)
         # EAGERLY process shared expert weights
         se._ensure_initialized()
-        # P1: Eagerly warmup fused SwiGLU compilation for SE (1-group)
-        if se._fused_swiglu:
-            from dsv4.ops.gemm_runner import warmup_fused_swiglu_compilation
-            K_packed = H // 2
-            N_packed_l1 = (2 * cfg.get("moe_intermediate_size", 3072)) // 2  # gate+up
-            warmup_fused_swiglu_compilation(
-                1, K_packed, N_packed_l1, dev,
-                swiglu_limit=cfg.get("swiglu_limit", 10.0),
-            )
-        # Fix activation global scales — _ensure_initialized sets gsa from l1_gs (which is 1.0)
-        # FIX: Same runtime gsa for SharedExpert
+        # BF16 fallback for shared expert — dequantize NVFP4 weights to BF16
         se._use_runtime_gsa = True
+        se.enable_bf16_fallback()  # sets _fused_swiglu=False, pre-materializes BF16 weights
         se_runners[li] = se
         if (li+1) % 10 == 0: print(f"  Built {li+1}/{n_layers} MoE layers")
         torch.cuda.empty_cache()