Fix NameError: dequantize_nvfp4 not in scope in forward_attention

The B3 fused q_a_norm path used dequantize_nvfp4 but it was only
imported in forward_layer, not forward_attention. Added local import.
This commit is contained in:
2026-06-02 21:52:29 +00:00
parent 9ba6476d3f
commit 9d4a014fad

View File

@@ -685,9 +685,9 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# With: fused rmsnorm+NVFP4 quantize → QuantizedActivation → q_b.run_from_quantized
# Saves: ~6 kernel launches per layer (rmsnorm 4+ + quantize 2 vs fused 2)
if q_norm_w is not None:
from dsv4.ops.quantize import rmsnorm_quantize_nvfp4 as _rmsnorm_quantize
from dsv4.ops.quantize import rmsnorm_quantize_nvfp4 as _rmsnorm_quantize, dequantize_nvfp4 as _dequantize_nvfp4
q_a_quant = _rmsnorm_quantize(q_a, q_norm_w.to(dev, torch.float32))
q_a = dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
q_a = _dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
_pt('q_b_start')
if q_norm_w is not None:
q = prod_lin['q_b'].run_from_quantized(q_a_quant)