[AMD] Fix DeepSeek-V4-Flash-FP8 on MI300 (#27972)

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py

@@ -46,12 +46,14 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.utils.common import is_gfx95_supported
+
 # OPUS gfx950 paged-prefill kernel is preferred when importable; otherwise fall
 # back to the Triton implementation below.
 try:
     from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_opus
 
-    _HAS_OPUS = True
+    _HAS_OPUS = is_gfx95_supported()
 except ImportError:
     pa_sparse_prefill_opus = None
     _HAS_OPUS = False

python/sglang/srt/mem_cache/memory_pool.py

@@ -863,7 +863,7 @@ class KVCache(abc.ABC):
         self.page_size = page_size
         self.dtype = dtype
         self.device = device
-        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn, torch.float8_e4m3fnuz):
             # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
             self.store_dtype = torch.uint8
         else: