test: quantize on different GPUs

test_quantize_gpu.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python3
+"""Test: quantize_activation_nvfp4 on different GPUs."""
+import torch
+from dsv4.ops.quantize import quantize_activation_nvfp4
+
+torch.manual_seed(42)
+
+for gpu in [0, 1]:
+    dev = f"cuda:{gpu}"
+    x = torch.randn(1, 7168, dtype=torch.bfloat16, device=dev) * 0.5
+    gsa = 0.000375
+    x_fp4, x_sf = quantize_activation_nvfp4(x, gsa)
+    has_nan = torch.isnan(x_fp4.view(torch.float16)).any().item() if x_fp4.dtype == torch.float4_e2m1fn_x2 else torch.isnan(x_fp4).any().item()
+    print(f"GPU {gpu} quantize: x_fp4 shape={x_fp4.shape} dtype={x_fp4.dtype} x_sf shape={x_sf.shape} has_nan={has_nan}")
+    print(f"  x_fp4 uint8 range: [{x_fp4.view(torch.uint8).min().item()}, {x_fp4.view(torch.uint8).max().item()}]")
+    print(f"  x_sf float range: [{x_sf.float().min().item():.6f}, {x_sf.float().max().item():.6f}]")