#!/usr/bin/env python3
"""Test shared expert on different GPUs."""
import torch
from dsv4.layers.shared_expert import Nvfp4SharedExpert
from dsv4.ops.quantize import quantize_weight_to_nvfp4

torch.manual_seed(42)

for gpu in [0, 1]:
    torch.cuda.set_device(gpu)
    dev = f"cuda:{gpu}"
    
    se = Nvfp4SharedExpert(hidden_size=7168, intermediate_size=3072, device=dev)
    
    # Create random BF16 weights and quantize to NVFP4
    gate_w = torch.randn(3072, 7168, dtype=torch.bfloat16, device=dev)
    up_w = torch.randn(3072, 7168, dtype=torch.bfloat16, device=dev)
    down_w = torch.randn(7168, 3072, dtype=torch.bfloat16, device=dev)
    
    gate_fp4, gate_sf, gate_gs = quantize_weight_to_nvfp4(gate_w)
    up_fp4, up_sf, up_gs = quantize_weight_to_nvfp4(up_w)
    down_fp4, down_sf, down_gs = quantize_weight_to_nvfp4(down_w)
    
    se.l1_fp4 = [torch.cat([gate_fp4, up_fp4], dim=0)]
    se.l1_sf = [torch.cat([gate_sf, up_sf], dim=0)]
    se.l1_gs = [1.0]
    se.l2_fp4 = [down_fp4]
    se.l2_sf = [down_sf]
    se.l2_gs = [1.0]
    
    # Input
    x = torch.randn(1, 7168, dtype=torch.bfloat16, device=dev)
    
    # Run
    out = se.run(x)
    has_nan = torch.isnan(out).any().item()
    print(f"GPU {gpu}: |out|={out.abs().max().item():.4f} has_nan={has_nan}")