From 16af96380f1cfa5378a9533008b430b9ada89e78 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 06:46:05 +0000 Subject: [PATCH] fix: use internal fields for Nvfp4Linear weight setup in test --- tests/unit/test_fused_router.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_fused_router.py b/tests/unit/test_fused_router.py index 74e3d502..4be3c1e3 100644 --- a/tests/unit/test_fused_router.py +++ b/tests/unit/test_fused_router.py @@ -35,13 +35,15 @@ def test_fused_router(): w_fp4, w_sf, ws2_val = quantize_weight_to_nvfp4(W_gate_bf16) # Build Nvfp4Linear for reference path + # The Nvfp4Linear expects stacked/contiguous weight tensors + from dsv4.layers.linear import Nvfp4Linear gate_lin = Nvfp4Linear(in_features=K, out_features=N, device=device) - gate_lin._weight_fp4 = [w_fp4.T.contiguous()] # K-major - gate_lin._weight_sf = [w_sf.T.contiguous()] + gate_lin._weight_fp4 = [w_fp4.contiguous()] + gate_lin._scale_b = [w_sf.contiguous()] gate_lin._gsb = [ws2_val] - gate_lin._activation_global_scale = None # will be set at runtime - gate_lin._ensure_stacked = lambda: None - gate_lin._ensure_initialized = lambda: None + gate_lin._activation_global_scale = None # set at runtime + gate_lin._ensure_stacked = lambda *a, **kw: None + gate_lin._ensure_initialized = lambda *a, **kw: None gate_lin.finalize_weights() # e_bias