From 6459fbca9a9c1f413dfdf5e1b2d81c17754a3bf1 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 05:41:33 +0000 Subject: [PATCH] fix: import forward_attention --- tests/unit/test_part_a_decode_diagnostics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_part_a_decode_diagnostics.py b/tests/unit/test_part_a_decode_diagnostics.py index 67769d8a..4b0e6109 100644 --- a/tests/unit/test_part_a_decode_diagnostics.py +++ b/tests/unit/test_part_a_decode_diagnostics.py @@ -42,7 +42,7 @@ def main(): from single_shot_inference import ( load_all_weights, make_nvfp4_linear, get_nvfp4_weight, rmsnorm, unweighted_rmsnorm, _apply_rope, build_rope_cache, - KVCache, Compressor, Indexer, forward_layer, moe_forward, + KVCache, Compressor, Indexer, forward_layer, forward_attention, moe_forward, _load_moe_weights_stacked, _load_shared_expert_weights, _cache_layer_weights_no_experts, )