test: add random FP4 data and kernel timeout

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