diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py index e384bc25e..cb35f5d0e 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 27a0d237f..85bd9a798 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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: