From 07168357cca0042b321c16aa56613e491e88c682 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 02:38:00 +0000 Subject: [PATCH] Fix o_a_proj weight loading: add BF16 fallback for grouped linear --- tests/unit/test_production_fmha_layer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/test_production_fmha_layer.py b/tests/unit/test_production_fmha_layer.py index 879182e6..a7c5744a 100644 --- a/tests/unit/test_production_fmha_layer.py +++ b/tests/unit/test_production_fmha_layer.py @@ -102,6 +102,10 @@ def main(): wo_a.load_nvfp4_weight(oa_w.to(dev), oa_ws.to(dev), oa_ws2.to(dev) if oa_ws2 is not None else None, oa_isc.to(dev) if oa_isc is not None else None) + else: + oa_bf = all_w.get(f"{pfx}.o_a_proj.weight") + if oa_bf is not None: + wo_a.set_bf16_weight(oa_bf.bfloat16().to(dev)) pl['o_a'] = wo_a; wo_a._use_runtime_gsa = True pl['o_b'] = make_nvfp4_linear(o_groups * o_rank, H, dev, all_w, pfx, 'o_b_proj') prod_lins[li] = pl