test: isolate fused SwiGLU — test no-clamp first
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user