[ROCm][CI] Fix AITER MLA fp8 decode metadata regression test (#48845)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Jul 21, 2026 at 19:19 UTC
61e10f0116ef45ea8c77271a9285185af87b0406
1 file changed
+5
-4
tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
+5
-4
@@ -40,10 +40,11 @@ BATCH_SIZE = 128
40
CONTEXT_LEN = 8192
41
PAGE_SIZE = 1
42
43
-# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
44
-# KV-cache -> fp8_e4m3 kv.
45
-EXPECTED_Q_DTYPE = torch.bfloat16
46
-EXPECTED_KV_DTYPE = torch.float8_e4m3fn
43
+# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
44
+# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
45
+# instead of hardcoding a literal (see #47276).
46
+EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
47
+EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
48
49
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
50
# is excluded: it holds raw device pointers, never equal across allocations.