Add check for float8 type in aiter mha

Change f8 kv-cache check in rocm_aiter_fa.py to account for both float8_e4m3fnuz
and float8_e4m3fn datatypes.

Signed-off-by: Doug Lehr <[email protected]>

Author: Doug Lehr

vllm/attention/ops/rocm_aiter_paged_attn.py

@@ -81,8 +81,8 @@ def forward_decode(
                 blocksparse_head_sliding_step=blocksparse_head_sliding_step)
 
         if "fp8" in kv_cache_dtype:
-            key_cache = key_cache.view(torch.float8_e4m3fnuz)
-            value_cache = value_cache.view(torch.float8_e4m3fnuz)
+            key_cache = key_cache.view(current_platform.fp8_dtype())
+            value_cache = value_cache.view(current_platform.fp8_dtype())
 
         if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
             # use blocksparse paged attention

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -479,8 +479,8 @@ def forward(
             )
 
         if self.kv_cache_dtype.startswith("fp8"):
-            key_cache = key_cache.view(torch.float8_e4m3fnuz)
-            value_cache = value_cache.view(torch.float8_e4m3fnuz)
+            key_cache = key_cache.view(current_platform.fp8_dtype())
+            value_cache = value_cache.view(current_platform.fp8_dtype())
 
         if not attn_metadata.use_cascade:
             cu_seqlens_q = attn_metadata.query_start_loc