biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: use internal fields for Nvfp4Linear weight setup in test

Browse Source
This commit is contained in:
biondizzle
2026-06-01 06:46:05 +00:00
parent 7f1f224c78
commit 16af96380f
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/unit/test_fused_router.py
View File

@@ -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