From 0c3796966dd103af5ee78bbec4379653540252be Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 13:19:27 +0000 Subject: [PATCH] =?UTF-8?q?Add=20BF16=20fallback=20for=20shared=20expert:?= =?UTF-8?q?=20dequantize=20NVFP4=20=E2=86=92=20BF16=20F.linear?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 41 ++++++++++++++++++++++++++++++++++++ single_shot_inference.py | 13 ++---------- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/dsv4/layers/shared_expert.py b/dsv4/layers/shared_expert.py index 8bbb034b..1c0fa724 100644 --- a/dsv4/layers/shared_expert.py +++ b/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: diff --git a/single_shot_inference.py b/single_shot_inference.py index 6f9ae641..b5e3a9a7 100644 --- a/single_shot_inference.py +++ b/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()