Fix fp8 amax in decode test

2026-05-19 15:28:17 +00:00
parent 4f23055450
commit 4b85605edf
1 changed files with 1 additions and 1 deletions
--- a/tests/test_decode_attention_b200.py
+++ b/tests/test_decode_attention_b200.py
@@ -292,7 +292,7 @@ def test_prefill_decode(layer_id, compress_ratio):
        kv_cache_bf16[:N_PREFILL] = kv_rope_prefill
        
        print(f"  Prefill: {N_PREFILL} tokens written to KV cache")
-        print(f"  KV cache fp8 amax: {kv_fp8_prefill.amax():.4f}")
+        print(f"  KV cache fp8 amax: {kv_fp8_prefill.float().abs().max():.4f}")
        print(f"  KV BF16 amax: {kv_rope_prefill.amax():.4f}")
        
        # Verify roundtrip: read back and compare