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

Add attention entropy diag (ATTN_DIAG), KV cache diag, --no-thinking mode

Browse Source
This commit is contained in:
biondizzle
2026-05-31 19:29:55 +00:00
parent 2a886fe0f2
commit a1b39adcaa
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
single_shot_inference.py
View File

@@ -66,6 +66,7 @@ INVERSE_ROPE = not _args.no_inverse_rope # If False, skip inverse RoPE on atten
SKIP_MHC = _args.skip_mhc # If True, bypass mHC and use simple residual connections (diagnostic)
MHC_DIAG = False # If True, print per-layer mHC diagnostics (B_l row/col sums, C_l values)
GROWTH_DIAG = True # If True, print per-layer residual growth analysis
ATTN_DIAG = True # If True, print per-layer attention entropy
# When True: applies inverse RoPE at query position → converts absolute→relative
# When False: leaves relative position encoding intact for output projection
# DSV4 partial RoPE only affects last 64/512 dims; first 448 are always un-RoPE'd
@@ -507,16 +508,16 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
attn_weights = torch.softmax(scores_raw.float(), dim=-1).to(torch.bfloat16)
attn_out = torch.matmul(attn_weights, v_expanded) # (n_h, T, hd)
attn_out = attn_out.permute(1, 0, 2) # (T, n_h, hd)
# Diagnostic: check attention entropy (how spread out the attention is)
if MHC_DIAG and li < 3:
# Diagnostic: check attention entropy and last-position weight
if ATTN_DIAG:
with torch.no_grad():
scores = torch.matmul(q_input, k_expanded.transpose(-1, -2)) * scale # (n_h, T, seq_len)
weights = torch.softmax(scores.float(), dim=-1) # (n_h, 1, seq_len)
# For head 0: what positions get the most weight?
w0 = weights[0, 0] # (seq_len,)
top3_pos = torch.topk(w0, min(3, seq_len))
entropy = -(w0 * (w0 + 1e-10).log()).sum().item()
print(f" L{li} attn: seq_len={seq_len} entropy={entropy:.2f} top3_pos={top3_pos.indices.tolist()} top3_w={top3_pos.values.tolist()}")
last_w = w0[-1].item() if seq_len > 0 else 0
print(f" L{li} attn: seq_len={seq_len} entropy={entropy:.2f} last_pos_w={last_w:.4f} top3_pos={top3_pos.indices.tolist()}")
else:
# Use FMHA kernel for longer sequences (padding effect is negligible)
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw