fix: import forward_attention

2026-06-03 05:41:33 +00:00
parent 91dfac34d8
commit 6459fbca9a
1 changed files with 1 additions and 1 deletions
--- 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,
    )