diff --git a/patches/test_nvfp4_mega_moe.py b/patches/test_nvfp4_mega_moe.py index c6ff286..da1d19c 100644 --- a/patches/test_nvfp4_mega_moe.py +++ b/patches/test_nvfp4_mega_moe.py @@ -78,26 +78,39 @@ def test_nvfp4_mega_moe(): import triton import triton.language as tl # Just manually pack a few tokens into FP4 for testing - # For now, zero-fill the symm buffer (the kernel should still launch without crash - # even with zero data — we just want to verify the kernel runs at all) - symm_buffer.x.zero_() - symm_buffer.x_sf.zero_() + # Write actual FP4 data to the activation buffer (random but valid packed E2M1) + symm_buffer.x[:num_tokens].copy_( + torch.randint(0, 256, (num_tokens, hidden // 2), dtype=torch.uint8, device=device).view(torch.int8)) + # Write valid UE4M3 scales (random but non-zero) + sf_data = torch.randn(num_tokens, hidden // 64, dtype=torch.float8_e4m3fn, device=device).abs().clamp(0.1, 5.0) + symm_buffer.x_sf[:num_tokens].copy_( + sf_data.view(torch.uint8).reshape(num_tokens, hidden // 64).contiguous().view(torch.int32)) # Write topk data directly - symm_buffer.topk_idx[:num_tokens].zero_() for i in range(num_tokens): for j in range(top_k): symm_buffer.topk_idx[i, j] = topk_ids[i, j].item() symm_buffer.topk_weights[i, j] = topk_weights[i, j].item() torch.cuda.synchronize() - print("Buffer populated (zeros for activations, real topk)") + print("Buffer populated with random FP4 data") # --- Run kernel --- y = torch.zeros(num_tokens, hidden, dtype=torch.bfloat16, device=device) print("Calling fp8_nvfp4_mega_moe...") try: fp8_nvfp4_mega_moe(y, l1_weights, l2_weights, symm_buffer) - torch.cuda.synchronize() - print(f"SUCCESS! y stats: min={y.min().item():.4f} max={y.max().item():.4f} mean={y.mean().item():.4f} nonzero={torch.count_nonzero(y).item()}") + # Use a sync with a manual timeout + done = torch.cuda.Event() + done.record() + import time + start = time.time() + while not done.query(): + if time.time() - start > 10: + print("TIMEOUT: kernel did not complete in 10s") + break + time.sleep(0.1) + else: + torch.cuda.synchronize() + print(f"SUCCESS! y stats: min={y.min().item():.4f} max={y.max().item():.4f} mean={y.mean().item():.4f} nonzero={torch.count_nonzero(y).item()}") except Exception as e: print(f"FAILED: {e}") raise