Add always-on alpha/x_sf debug prints for L1 and L2 GEMM calls

This commit is contained in:
2026-05-15 03:59:07 +00:00
parent 9c318c3353
commit 9975558c23

View File

@@ -308,6 +308,11 @@ def nvfp4_mega_moe_full(
topk_ids = symm_buffer.topk_idx[:num_tokens]
topk_weights = symm_buffer.topk_weights[:num_tokens]
# ALWAYS-ON debug: alpha and scale ranges
_x_sf_f32 = x_sf.to(torch.float32)
_igs = l1_global_scale if isinstance(l1_global_scale, float) else l1_global_scale.item() if hasattr(l1_global_scale, 'item') else float(l1_global_scale)
print(f"[ALPHA L1] alpha={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.max().item():.4e}] x_fp4_absmax={x_fp4.view(torch.int8).abs().max().item()}")
# Convert global expert IDs to local expert IDs.
# vLLM's symm_buffer stores global IDs (0..383) but our weight tensors
# are indexed by local ID (0..47). Each rank handles a contiguous chunk:
@@ -356,6 +361,11 @@ def nvfp4_mega_moe_full(
# Step 4: Quantize L1 output → FP4
l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated)
# ALWAYS-ON debug: L2 alpha and scale ranges
_l1sf_f32 = l1_sf_out.to(torch.float32)
_l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item() if hasattr(l2_global_scale, 'item') else float(l2_global_scale)
print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}] activated amax={activated.abs().max().item():.4e}")
# Step 5: L2 GEMM (native NVFP4 block-scaled MMA)
l2_output = nvfp4_mega_moe_l2(
l1_fp4, l1_sf_out, l2_w, l2_sf,