[Perf] Improve FP8 dispatch storage release

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: adam-xiaoyao

sonicmoe/ernie_compat/mlp_node_v2.py

@@ -504,22 +504,10 @@ def _attach_fp8_weight_attrs(
         w2: torch.Tensor,
         payload: dict[str, tuple[torch.Tensor, torch.Tensor]],
     ) -> None:
-        w1.fp8_weight = {
-            "w1_fused": payload["w1_fused"][0],
-            "w1T_varlen": payload["w1T_varlen"][0],
-        }
-        w1.fp8_scale = {
-            "w1_fused": payload["w1_fused"][1],
-            "w1T_varlen": payload["w1T_varlen"][1],
-        }
-        w2.fp8_weight = {
-            "w2_varlen": payload["w2_varlen"][0],
-            "w2_dgated": payload["w2_dgated"][0],
-        }
-        w2.fp8_scale = {
-            "w2_varlen": payload["w2_varlen"][1],
-            "w2_dgated": payload["w2_dgated"][1],
-        }
+        w1.fp8 = payload["w1_fused"]
+        w1.transposed_fp8 = payload["w1T_varlen"]
+        w2.fp8 = payload["w2_varlen"]
+        w2.transposed_fp8 = payload["w2_dgated"]
 
     # ── Weight layout helpers (instance-scoped) ─────────────────────────────
 
@@ -907,7 +895,6 @@ def forward(
                     False,                          # is_inference_mode_enabled
                     False,                          # use_low_precision_postact_buffer
                     fp8_activation_payload,
-                    fp8_weight_payload,
                 )
 
             # ── DownProjection forward (via FakeCtx) ─────────────────────────
@@ -924,7 +911,6 @@ def forward(
                     activation_type,
                     None,                           # fp8_protocol
                     None,
-                    fp8_weight_payload,
                     router_scores_expert_order,
                     router_scores_token_order,
                     score_src_idx,

sonicmoe/functional/__init__.py

@@ -1283,12 +1283,19 @@ def forward(
                     if prequant_activation_payload is not None:
                         x_fp8_pre, x_scales_pre = prequant_activation_payload
                     w1_fused_payload = _get_fp8_weight_attr(w1, "fp8")
-                    z, y1 = _fused_blockscaled_gated_forward(
-                        x, w1, expert_frequency_offset, x_gather_idx,
-                        w1_fp8_pre=w1_fused_payload,
-                        x_fp8_pre=x_fp8_pre, x_scales_pre=x_scales_pre,
-                        store_z=not cfg.recompute_z,
-                        fuse_y1_quant=fuse_y1,
+                    fuse_y1 = _use_fuse_y1_quant()
+                    z, y1, y1_fp8_fused, y1_scales_fused = (
+                        _fused_blockscaled_gated_forward(
+                            x,
+                            w1,
+                            expert_frequency_offset,
+                            x_gather_idx,
+                            w1_fp8_pre=w1_fused_payload,
+                            x_fp8_pre=x_fp8_pre,
+                            x_scales_pre=x_scales_pre,
+                            store_z=not cfg.recompute_z,
+                            fuse_y1_quant=fuse_y1,
+                        )
                     )
                     if cfg.recompute_z:
                         # Forward skips preact storage; z is materialized just-in-time
@@ -2053,9 +2060,22 @@ def forward(
 
         # w2 decoupling: in FP8+aligned+fused_gated mode, backward doesn't
         # read bf16 w2 data (uses fp8 dgated cache + metadata).  This enables
-        # stash_bf16_to_cpu() to resize_(0) the bf16 param storage.
-        _w2_decouple = z_is_fp8 and cfg.fused_gated
+        # clear_param_storage("moe_expert") to release bf16 expert weight storage
+        # without requiring recompute_z/save_z_fp8.
+        _w2_decouple = (
+            cfg.enabled
+            and use_quack_gemm
+            and cfg.alignment_assumed
+            and cfg.fused_gated
+        )
         ctx._w2_decoupled = _w2_decouple
+        if _w2_decouple:
+            ctx._w2_dgated_fp8, ctx._w2_dgated_scales = _get_fp8_weight_attr(
+                w2, "transposed_fp8"
+            )
+            ctx._w2_shape = w2.shape  # (H, I, E)
+            ctx._w2_dtype = w2.dtype
+            ctx._w2_device = w2.device
 
         if z_is_fp8:
             recompute_args = _PREQUANTIZED_SCALES.pop("z_fp8_recompute", None)
@@ -2099,10 +2119,6 @@ def forward(
                     else:
                         z_fp8, z_raw_scales = quantize_activation_blockscaled_fast(z)
             if _w2_decouple:
-                ctx._w2_dgated_fp8, ctx._w2_dgated_scales = _get_fp8_weight_attr(w2, "transposed_fp8")
-                ctx._w2_shape = w2.shape  # (H, I, E)
-                ctx._w2_dtype = w2.dtype
-                ctx._w2_device = w2.device
                 ctx.save_for_backward(
                     z_fp8,
                     z_raw_scales,
@@ -2129,17 +2145,30 @@ def forward(
                     s_reverse_scatter_idx,
                 )
         else:
-            ctx.save_for_backward(
-                z,
-                w2,
-                b2,
-                topk_scores,
-                topk_scores_expert_order,
-                expert_frequency_offset,
-                x_gather_idx,
-                s_scatter_idx,
-                s_reverse_scatter_idx,
-            )
+            if _w2_decouple:
+                ctx.save_for_backward(
+                    z,
+                    # w2 omitted — backward uses ctx._w2_dgated_fp8 + metadata
+                    b2,
+                    topk_scores,
+                    topk_scores_expert_order,
+                    expert_frequency_offset,
+                    x_gather_idx,
+                    s_scatter_idx,
+                    s_reverse_scatter_idx,
+                )
+            else:
+                ctx.save_for_backward(
+                    z,
+                    w2,
+                    b2,
+                    topk_scores,
+                    topk_scores_expert_order,
+                    expert_frequency_offset,
+                    x_gather_idx,
+                    s_scatter_idx,
+                    s_reverse_scatter_idx,
+                )
 
         # Keep w2 FP8 cache — backward hits cache (~38µs savings) at ~37MB memory cost.
         # The cache auto-invalidates via w._version when optimizer updates weights.
@@ -2201,20 +2230,35 @@ def backward(ctx, dout: torch.Tensor):
             # Defer dequantize: FP8 path uses fused kernel, others lazy-dequant
             z = None
         else:
-            (
-                z,
-                w2,
-                b2,
-                topk_scores,
-                topk_scores_expert_order,
-                expert_frequency_offset,
-                x_gather_idx,
-                s_scatter_idx,
-                s_reverse_scatter_idx,
-            ) = ctx.saved_tensor()
-            w2_shape = w2.shape
-            w2_dtype = w2.dtype
-            w2_device = w2.device
+            if ctx._w2_decoupled:
+                (
+                    z,
+                    b2,
+                    topk_scores,
+                    topk_scores_expert_order,
+                    expert_frequency_offset,
+                    x_gather_idx,
+                    s_scatter_idx,
+                    s_reverse_scatter_idx,
+                ) = ctx.saved_tensor()
+                w2_shape = ctx._w2_shape
+                w2_dtype = ctx._w2_dtype
+                w2_device = ctx._w2_device
+            else:
+                (
+                    z,
+                    w2,
+                    b2,
+                    topk_scores,
+                    topk_scores_expert_order,
+                    expert_frequency_offset,
+                    x_gather_idx,
+                    s_scatter_idx,
+                    s_reverse_scatter_idx,
+                ) = ctx.saved_tensor()
+                w2_shape = w2.shape
+                w2_dtype = w2.dtype
+                w2_device = w2.device
             z_fp8 = z_raw_scales_u8 = None
             if getattr(ctx, "_needs_z_recompute_bf16", False):
                 # Replace the zero-storage bf16 placeholder with a freshly recomputed