Skip to content

[Perf] Improve FP8 dispatch storage release #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
A-nnonymous wants to merge 2 commits into
base: paddle
Choose a base branch
from
Open

[Perf] Improve FP8 dispatch storage release #43

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2b800cd
[Perf] Improve FP8 dispatch storage release
A-nnonymous Jun 10, 2026
d6222b4
[BugFix] Restore router score gradients
A-nnonymous Jun 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 4 additions & 18 deletions sonicmoe/ernie_compat/mlp_node_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) ─────────────────────────────

Expand Down Expand Up @@ -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) ─────────────────────────
Expand All @@ -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,
Expand Down
123 changes: 82 additions & 41 deletions sonicmoe/functional/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -2024,10 +2031,7 @@ def forward(
# ctx.needs_input_grad. Defaulting to True is safe: if the caller truly
# doesn't need ds, the autograd engine simply discards it.

if not hasattr(topk_scores, "stop_gradient"):
ctx._topk_scores_needs_grad = False
else:
ctx._topk_scores_needs_grad = not topk_scores.stop_gradient
ctx._topk_scores_needs_grad = True
ctx._fp8_combine_grad_handle = fp8_combine_grad_handle
ctx._has_topk_scores_expert_order = topk_scores_expert_order is not None
ctx._has_topk_scores_token_order = topk_scores_token_order is not None
Expand All @@ -2053,9 +2057,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)
Expand Down Expand Up @@ -2099,10 +2116,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,
Expand All @@ -2129,17 +2142,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.
Expand Down Expand Up @@ -2201,20 +2227,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
Expand Down