test: isolate fused SwiGLU — test no-clamp first

This commit is contained in:
2026-06-02 08:23:28 +00:00
parent ae4506d722
commit 53362d2579

View File

@@ -92,8 +92,30 @@ def test_fused_swiglu_compilation():
x_padded = torch.zeros(128, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2)
x_padded.view(torch.uint8)[:tokens] = x_fp4.view(torch.uint8)
# Run FUSED path
print(" Running fused SwiGLU GEMM...")
# Test WITHOUT clamp first (swiglu_limit=0) to isolate SiLU-only path
print("\n--- Smoke test: fused kernel (no clamp) ---")
swiglu_limit_test = 0.0
warmup_fused_swiglu_compilation(
num_experts, K_packed, N_packed, device,
swiglu_limit=swiglu_limit_test,
)
l1_fused_noclamp = run_fused_swiglu_grouped_gemm(
mat_a=x_padded, mat_b=mat_b,
scale_a=scale_a, scale_b=scale_b,
expert_offsets=padded_offsets,
global_scale_a=gsa, global_scale_b=gsb,
swiglu_limit=swiglu_limit_test,
)[:tokens]
has_nan_nc = torch.isnan(l1_fused_noclamp).any().item()
max_val_nc = l1_fused_noclamp.abs().max().item()
print(f" No-clamp output: NaN={has_nan_nc} max={max_val_nc:.4f}")
if has_nan_nc or max_val_nc > 1e6:
print(" ❌ No-clamp path has issues")
return False
print(" ✅ No-clamp path works")
# Test WITH clamp
print("\n--- Smoke test: fused kernel (with clamp) ---")
l1_fused = run_fused_swiglu_grouped_gemm(
mat_a=x_padded, mat_b=mat_b,
scale_a=scale_a, scale_b=scale_b,