biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add --skip-mhc flag for simple residual diagnostic

Browse Source 
When enabled, bypasses mHC pre/post blocks and uses direct residual
connections with 0.1 scaling. This helps isolate whether the mHC
implementation is causing the garbage output.
This commit is contained in:
biondizzle
2026-05-31 11:33:41 +00:00
parent d201a9334e
commit 631e6ea3e4
1 changed files with 25 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
single_shot_inference.py
View File

@@ -48,6 +48,7 @@ def parse_args():
p = argparse.ArgumentParser(description='DSV4 Single-Shot Inference')
p.add_argument('--no-inverse-rope', action='store_true', help='Skip inverse RoPE on attention output')
p.add_argument('--skip-moe', action='store_true', help='Only use shared expert (skip routed)')
p.add_argument('--skip-mhc', action='store_true', help='Bypass mHC, use simple residual (diagnostic)')
p.add_argument('--max-tokens', type=int, default=512, help='Max new tokens to generate')
p.add_argument('--prompt', type=str, default=None, help='Override prompt')
return p.parse_args()
@@ -61,7 +62,7 @@ PROMPT = _args.prompt or "The capital of France is"
NUM_GPUS = 8
SKIP_ROUTED_MOE = _args.skip_moe # If True, only use shared expert (debug)
INVERSE_ROPE = not _args.no_inverse_rope # If False, skip inverse RoPE on attention output (diagnostic)
LAYER_TRACE = False # If True, trace every computation step for first layer of first prefill token
SKIP_MHC = _args.skip_mhc # If True, bypass mHC and use simple residual connections (diagnostic)
MHC_DIAG = True # If True, print per-layer mHC diagnostics (B_l row/col sums, C_l values)
# When True: applies inverse RoPE at query position → converts absolute→relative
# When False: leaves relative position encoding intact for output projection
@@ -398,8 +399,13 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
# ATTENTION SUB-BLOCK
# ==================================================================
# -- mHC pre_block (attention) --
x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H)
if SKIP_MHC:
# Simple residual: skip mHC, use direct input
x_in = X_l[:, 0, :] # Just take stream 0
attn_ctx = None
else:
# -- mHC pre_block (attention) --
x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H)
if MHC_DIAG: # mHC diagnostics
B_l, C_l = attn_ctx.B_l, attn_ctx.C_l
print(f" L{li} pre_attn: |X_l|={X_l.abs().max().item():.2f} |x_in|={x_in.abs().max().item():.2f}", flush=True)
@@ -546,8 +552,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
w[f"{pre}.o_b_proj.weight_scale"],
w[f"{pre}.o_b_proj.weight_scale_2"]) # (T, H)
# -- mHC post_block (attention) --
X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H)
if SKIP_MHC:
X_mid = X_l[:, 0, :].unsqueeze(1).expand(-1, 4, -1) + F_attn.unsqueeze(1) * 0.1
else:
# -- mHC post_block (attention) --
X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H)
# Diagnostic: check mHC is stabilizing the residual
if MHC_DIAG: # mHC diagnostics
B_l, C_l = attn_ctx.B_l, attn_ctx.C_l
@@ -562,8 +571,12 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
# FFN SUB-BLOCK
# ==================================================================
# -- mHC pre_block (FFN) --
x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) # (T, H)
if SKIP_MHC:
x_ffn = X_mid[:, 0, :] # Just take stream 0
ffn_ctx = None
else:
# -- mHC pre_block (FFN) --
x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) # (T, H)
# -- RMSNorm (pre-norm before FFN) --
x_ffn_normed = ffn_norm.forward(x_ffn) # (T, H) BF16
@@ -571,8 +584,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
# -- MoE + shared expert --
F_ffn = moe_forward(x_ffn_normed, w, li, cfg, token_id, device)
# -- mHC post_block (FFN) --
X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H)
if SKIP_MHC:
X_next = X_mid + F_ffn.unsqueeze(1) * 0.1
else:
# -- mHC post_block (FFN) --
X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H)
if MHC_DIAG: # ffn mHC diagnostics
B_l_ffn, C_l_ffn = ffn_ctx.B_l, ffn_ctx.C_l
print(f" L{li} ffn: |X_mid|={X_mid.abs().max().item():.2f} |F_ffn|={F_ffn.abs().max().item():.2f} |B|={B_l_ffn.abs().max().item():.4f} |C|={C_l_ffn.abs().max().item():.4f} |X_next|={X_next.abs().max().item():.2f}", flush=True)